| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """RLHF with FSDP2 training and vLLM tensor-parallel inference using **NCCL M2N** |
| 4 | sharding-aware weight transfer. |
| 5 | |
| 6 | Layout (4 GPUs, no colocation): |
| 7 | * GPUs 0-1: two FSDP2 training workers, one per GPU. |
| 8 | * GPUs 2-3: one vLLM ``LLM`` actor with ``tensor_parallel_size=2``. |
| 9 | |
| 10 | The trainer and the inference workers share one NCCL communicator of 4 ranks: |
| 11 | trainer ranks ``[0, 2)`` and inference ranks ``[2, 4)``. Every parameter moves |
| 12 | with a single ``reshard`` that redistributes it from the FSDP layout to the |
| 13 | inference layout — the trainer sends its local shards and never all-gathers a |
| 14 | full tensor, which broadcast-based weight sync would force it to do. |
| 15 | |
| 16 | Every FSDP rank builds an ``M2NTrainerWeightTransferEngine`` and calls |
| 17 | ``send_weights()``; all ranks run every reshard, and only rank 0 drives the |
| 18 | inference side through its ``RayVLLMWeightSyncClient``. |
| 19 | |
| 20 | Requires the ``nccl-extensions`` package (NCCL 2.30.5+) and a |
| 21 | ``VLLM_NCCL_SO_PATH`` pointing at the same ``libnccl.so`` that |
| 22 | ``libnccl_m2n.so`` was linked against. |
| 23 | |
| 24 | This example was written for 4xH100. |
| 25 | """ |
| 26 | |
| 27 | from __future__ import annotations |
| 28 | |
| 29 | import os |
| 30 | |
| 31 | import pybase64 as base64 |
| 32 | import ray |
| 33 | import torch |
| 34 | import torch.distributed as dist |
| 35 | from huggingface_hub import snapshot_download |
| 36 | from torch.distributed.fsdp import fully_shard |
| 37 | from transformers import AutoModelForCausalLM |
| 38 | |
| 39 | from vllm import LLM, SamplingParams |
| 40 | from vllm.config import WeightTransferConfig |
| 41 | from vllm.distributed.device_communicators.pynccl_wrapper import NCCLLibrary |
| 42 | from vllm.distributed.weight_transfer import ( |
| 43 | RayVLLMWeightSyncClient, |
| 44 | WeightTransferTrainerFactory, |
| 45 | ) |
| 46 | from vllm.distributed.weight_transfer.m2n_source import DTensorModuleSource |
| 47 | from vllm.distributed.weight_transfer.m2n_trainer import M2NTrainerInitInfo |
| 48 | from vllm.utils.network_utils import get_ip, get_open_port |
| 49 | |
| 50 | MODEL_NAME = "Qwen/Qwen3-0.6B" |
| 51 | |
| 52 | FSDP_WORLD_SIZE = 2 |
| 53 | INFERENCE_TP_SIZE = 2 |
| 54 | |
| 55 | |
| 56 | class MyLLM(LLM): |
| 57 | """LLM subclass that keeps Ray from pinning it to a single device.""" |
| 58 | |
| 59 | def __init__(self, *args, **kwargs): |
| 60 | os.environ.pop("CUDA_VISIBLE_DEVICES", None) |
| 61 | super().__init__(*args, **kwargs) |
| 62 | |
| 63 | def ready(self): |
| 64 | return True |
| 65 | |
| 66 | |
| 67 | @ray.remote(num_cpus=1, num_gpus=1) |
| 68 | class FSDPTrainWorker: |
| 69 | """One FSDP2 worker per GPU. Rank 0 is the weight-transfer sender.""" |
| 70 | |
| 71 | def __init__( |
| 72 | self, |
| 73 | model_name: str, |
| 74 | rank: int, |
| 75 | fsdp_master_addr: str, |
| 76 | fsdp_master_port: int, |
| 77 | ): |
| 78 | self.rank = rank |
| 79 | |
| 80 | os.environ["MASTER_ADDR"] = fsdp_master_addr |
| 81 | os.environ["MASTER_PORT"] = str(fsdp_master_port) |
| 82 | |
| 83 | dist.init_process_group(backend="nccl", rank=rank, world_size=FSDP_WORLD_SIZE) |
| 84 | torch.accelerator.set_device_index(0) |
| 85 | |
| 86 | model = AutoModelForCausalLM.from_pretrained( |
| 87 | model_name, torch_dtype=torch.bfloat16 |
| 88 | ).cuda() |
| 89 | for layer in model.model.layers: |
| 90 | fully_shard(layer) |
| 91 | fully_shard(model) |
| 92 | self.model = model |
| 93 | |
| 94 | def ready(self): |
| 95 | return True |
| 96 | |
| 97 | def setup_engine(self, llm_handle, nccl_unique_id_b64, world_size, num_workers): |
| 98 | """Build the trainer engine on every FSDP rank. |
| 99 | |
| 100 | `DTensorModuleSource` reads each parameter's FSDP device mesh and |
| 101 | placements, so the engine knows the source layout without gathering |
| 102 | anything. Rank 0 additionally drives the inference-side handshake. |
| 103 | """ |
| 104 | self.engine = WeightTransferTrainerFactory.trainer_init( |
| 105 | init_info=M2NTrainerInitInfo( |
| 106 | nccl_unique_id_b64=nccl_unique_id_b64, |
| 107 | world_size=world_size, |
| 108 | num_trainer_ranks=FSDP_WORLD_SIZE, |
| 109 | # One DP group of TP=2 workers; declared so both sides |
| 110 | # describe the destination the same way. |
| 111 | dst_mesh_dims=(num_workers // INFERENCE_TP_SIZE, INFERENCE_TP_SIZE), |
| 112 | rank=self.rank, # FSDP rank; sender is 0 |
| 113 | ), |
| 114 | client=RayVLLMWeightSyncClient(llm_handle), |
| 115 | source=DTensorModuleSource(self.model, FSDP_WORLD_SIZE), |
| 116 | ) |
| 117 | |
| 118 | def send_weights(self): |
| 119 | """Called on all ranks concurrently; every rank runs every reshard.""" |
| 120 | self.engine.send_weights() |
| 121 | |
| 122 | |
| 123 | def main(): |
| 124 | ray.init() |
| 125 | |
| 126 | local_model_path = snapshot_download(MODEL_NAME) |
| 127 | print(f"[init] Model downloaded to {local_model_path}") |
| 128 | |
| 129 | fsdp_master_addr = get_ip() |
| 130 | fsdp_master_port = get_open_port() |
| 131 | |
| 132 | fsdp_workers = [ |
| 133 | FSDPTrainWorker.remote( |
| 134 | local_model_path, rank, fsdp_master_addr, fsdp_master_port |
| 135 | ) |
| 136 | for rank in range(FSDP_WORLD_SIZE) |
| 137 | ] |
| 138 | ray.get([w.ready.remote() for w in fsdp_workers]) |
| 139 | print(f"[init] {FSDP_WORLD_SIZE} FSDP workers ready.") |
| 140 | |
| 141 | llm = ray.remote(num_cpus=0, num_gpus=0)(MyLLM).remote( |
| 142 | model=local_model_path, |
| 143 | enforce_eager=True, |
| 144 | tensor_parallel_size=INFERENCE_TP_SIZE, |
| 145 | distributed_executor_backend="ray", |
| 146 | weight_transfer_config=WeightTransferConfig(backend="nccl_m2n"), |
| 147 | load_format="dummy", |
| 148 | ) |
| 149 | ray.get(llm.ready.remote()) |
| 150 | num_workers = ray.get(llm.get_world_size.remote()) |
| 151 | print(f"[init] vLLM ready with {num_workers} inference workers.") |
| 152 | |
| 153 | prompts = [ |
| 154 | "Hello, my name is", |
| 155 | "The capital of France is", |
| 156 | ] |
| 157 | sampling_params = SamplingParams(temperature=0) |
| 158 | |
| 159 | outputs = ray.get(llm.generate.remote(prompts, sampling_params)) |
| 160 | print("-" * 60) |
| 161 | print("BEFORE weight sync (dummy weights):") |
| 162 | for output in outputs: |
| 163 | print(f"Prompt: {output.prompt!r}") |
| 164 | print(f"Generated: {output.outputs[0].text!r}") |
| 165 | print("-" * 60) |
| 166 | |
| 167 | # Trainer ranks [0, FSDP_WORLD_SIZE) and inference ranks after them share |
| 168 | # one communicator, so the two meshes are contiguous rank intervals. |
| 169 | nccl = NCCLLibrary() |
| 170 | nccl_unique_id_b64 = base64.b64encode( |
| 171 | bytes(nccl.ncclGetUniqueId().internal) |
| 172 | ).decode() |
| 173 | world_size = FSDP_WORLD_SIZE + num_workers |
| 174 | |
| 175 | print("[transfer] Initializing nccl_m2n weight transfer (all FSDP ranks)...") |
| 176 | ray.get( |
| 177 | [ |
| 178 | w.setup_engine.remote(llm, nccl_unique_id_b64, world_size, num_workers) |
| 179 | for w in fsdp_workers |
| 180 | ] |
| 181 | ) |
| 182 | |
| 183 | print("[sync] Resharding FSDP -> vLLM...") |
| 184 | ray.get([w.send_weights.remote() for w in fsdp_workers]) |
| 185 | print("[sync] Weight transfer complete.") |
| 186 | |
| 187 | outputs_updated = ray.get(llm.generate.remote(prompts, sampling_params)) |
| 188 | print("-" * 60) |
| 189 | print("AFTER weight sync (real weights):") |
| 190 | for output in outputs_updated: |
| 191 | print(f"Prompt: {output.prompt!r}") |
| 192 | print(f"Generated: {output.outputs[0].text!r}") |
| 193 | print("-" * 60) |
| 194 | |
| 195 | |
| 196 | if __name__ == "__main__": |
| 197 | main() |