| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Demonstrate checkpoint-coordinate sparse NCCL updates with expert parallelism. |
| 4 | |
| 5 | The trainer and vLLM start from the same Qwen3 MoE checkpoint. The trainer |
| 6 | modifies rows from two global experts, converts its fused expert storage back to |
| 7 | per-expert checkpoint coordinates, and sends both patches through one |
| 8 | ``send_weights()`` lifecycle. Every vLLM rank receives both patches; the native |
| 9 | loader applies its local expert and skips the foreign expert. |
| 10 | |
| 11 | This example uses three GPUs on one node: one for the full BF16 Hugging Face |
| 12 | trainer model and two for a TP2/EP2 vLLM inference engine. Unspecified checkpoint |
| 13 | elements keep their initialized values, so sparse updates require a known shared |
| 14 | baseline. ``SPARSE_NCCL_MODEL`` may point to a compatible local Qwen3 MoE |
| 15 | checkpoint. |
| 16 | """ |
| 17 | |
| 18 | import os |
| 19 | from contextlib import suppress |
| 20 | |
| 21 | import ray |
| 22 | import torch |
| 23 | from ray.util.placement_group import placement_group |
| 24 | from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy |
| 25 | from transformers import AutoModelForCausalLM |
| 26 | |
| 27 | from vllm import LLM, SamplingParams |
| 28 | from vllm.config import WeightTransferConfig |
| 29 | from vllm.distributed.weight_transfer import ( |
| 30 | RayVLLMWeightSyncClient, |
| 31 | WeightTransferTrainerFactory, |
| 32 | ) |
| 33 | from vllm.distributed.weight_transfer.sparse_nccl_engine import ( |
| 34 | SparseNCCLTrainerInitInfo, |
| 35 | SparseWeightPatch, |
| 36 | ) |
| 37 | from vllm.utils.network_utils import get_ip, get_open_port |
| 38 | |
| 39 | MODEL_NAME = os.environ.get("SPARSE_NCCL_MODEL", "Qwen/Qwen3-30B-A3B") |
| 40 | INFERENCE_TP_SIZE = 2 |
| 41 | PATCHED_LAYER = 0 |
| 42 | PATCHED_ROWS = 2 |
| 43 | PROMPTS = [ |
| 44 | "Hello, my name is", |
| 45 | "The president of the United States is", |
| 46 | "The capital of France is", |
| 47 | "The future of AI is", |
| 48 | ] |
| 49 | SAMPLING_PARAMS = SamplingParams(temperature=0.0, max_tokens=1) |
| 50 | |
| 51 | |
| 52 | class MyLLM(LLM): |
| 53 | """Configure the vLLM worker for Ray placement group execution.""" |
| 54 | |
| 55 | def __init__(self, *args, **kwargs): |
| 56 | os.environ["VLLM_RAY_BUNDLE_INDICES"] = ",".join( |
| 57 | str(index) for index in range(INFERENCE_TP_SIZE) |
| 58 | ) |
| 59 | super().__init__(*args, **kwargs) |
| 60 | |
| 61 | |
| 62 | @ray.remote(num_gpus=1) |
| 63 | class TrainModel: |
| 64 | """Own the trainer model and sparse NCCL sender on one GPU.""" |
| 65 | |
| 66 | def __init__(self, model_name: str): |
| 67 | self.model = AutoModelForCausalLM.from_pretrained( |
| 68 | model_name, |
| 69 | torch_dtype=torch.bfloat16, |
| 70 | ).to("cuda:0") |
| 71 | self.model.eval() |
| 72 | |
| 73 | config = self.model.config |
| 74 | if config.model_type != "qwen3_moe" or config.num_experts % INFERENCE_TP_SIZE: |
| 75 | raise RuntimeError( |
| 76 | "This recipe requires a Qwen3 MoE model whose experts divide " |
| 77 | f"evenly across TP{INFERENCE_TP_SIZE}" |
| 78 | ) |
| 79 | self.expert_ids = (0, config.num_experts // 2) |
| 80 | self.expert_intermediate_size = config.moe_intermediate_size |
| 81 | fused_name = f"model.layers.{PATCHED_LAYER}.mlp.experts.gate_up_proj" |
| 82 | try: |
| 83 | self.gate_up_proj = self.model.get_parameter(fused_name) |
| 84 | except AttributeError as exc: |
| 85 | raise RuntimeError( |
| 86 | f"Expected trainer model to expose `{fused_name}`" |
| 87 | ) from exc |
| 88 | |
| 89 | expected_shape = ( |
| 90 | config.num_experts, |
| 91 | 2 * self.expert_intermediate_size, |
| 92 | config.hidden_size, |
| 93 | ) |
| 94 | if self.gate_up_proj.shape != expected_shape: |
| 95 | raise RuntimeError( |
| 96 | f"Unexpected fused expert shape: {self.gate_up_proj.shape} " |
| 97 | f"!= {expected_shape}" |
| 98 | ) |
| 99 | |
| 100 | self.master_address = get_ip() |
| 101 | self.master_port = get_open_port() |
| 102 | self.engine = None |
| 103 | |
| 104 | def init_sparse_engine(self, world_size: int, llm_handle) -> None: |
| 105 | self.engine = WeightTransferTrainerFactory.trainer_init( |
| 106 | init_info=SparseNCCLTrainerInitInfo( |
| 107 | master_address=self.master_address, |
| 108 | master_port=self.master_port, |
| 109 | world_size=world_size, |
| 110 | rank=0, |
| 111 | ), |
| 112 | client=RayVLLMWeightSyncClient(llm_handle), |
| 113 | ) |
| 114 | |
| 115 | @torch.no_grad() |
| 116 | def patch_and_send(self) -> tuple[list[str], int]: |
| 117 | if self.engine is None: |
| 118 | raise RuntimeError("Sparse NCCL engine is not initialized") |
| 119 | |
| 120 | patches = [] |
| 121 | for expert_id in self.expert_ids: |
| 122 | checkpoint_weight = self.gate_up_proj[ |
| 123 | expert_id, : self.expert_intermediate_size |
| 124 | ] |
| 125 | original_rows = checkpoint_weight[:PATCHED_ROWS].clone() |
| 126 | replacement_rows = original_rows.flip(0) |
| 127 | if torch.equal(original_rows, replacement_rows): |
| 128 | raise RuntimeError(f"Expert {expert_id} patch would be a no-op") |
| 129 | checkpoint_weight[:PATCHED_ROWS] = replacement_rows |
| 130 | |
| 131 | hidden_size = checkpoint_weight.shape[1] |
| 132 | flat_indices = torch.arange( |
| 133 | PATCHED_ROWS * hidden_size, |
| 134 | device=checkpoint_weight.device, |
| 135 | dtype=torch.int32, |
| 136 | ) |
| 137 | patches.append( |
| 138 | SparseWeightPatch( |
| 139 | name=( |
| 140 | f"model.layers.{PATCHED_LAYER}.mlp.experts." |
| 141 | f"{expert_id}.gate_proj.weight" |
| 142 | ), |
| 143 | full_shape=tuple(checkpoint_weight.shape), |
| 144 | indices=flat_indices, |
| 145 | values=replacement_rows.reshape(-1).contiguous(), |
| 146 | ) |
| 147 | ) |
| 148 | |
| 149 | self.engine.send_weights(patches) |
| 150 | return [patch.name for patch in patches], sum( |
| 151 | patch.indices.numel() for patch in patches |
| 152 | ) |
| 153 | |
| 154 | def shutdown_engine(self) -> None: |
| 155 | if self.engine is not None: |
| 156 | self.engine.shutdown() |
| 157 | self.engine = None |
| 158 | |
| 159 | |
| 160 | def launch_llm(scheduling_strategy: PlacementGroupSchedulingStrategy): |
| 161 | return ray.remote( |
| 162 | num_cpus=0, |
| 163 | num_gpus=0, |
| 164 | scheduling_strategy=scheduling_strategy, |
| 165 | )(MyLLM).remote( |
| 166 | model=MODEL_NAME, |
| 167 | enforce_eager=True, |
| 168 | tensor_parallel_size=INFERENCE_TP_SIZE, |
| 169 | enable_expert_parallel=True, |
| 170 | expert_placement_strategy="linear", |
| 171 | moe_backend="triton", |
| 172 | distributed_executor_backend="ray", |
| 173 | gpu_memory_utilization=0.7, |
| 174 | enable_prefix_caching=False, |
| 175 | weight_transfer_config=WeightTransferConfig(backend="sparse_nccl"), |
| 176 | ) |
| 177 | |
| 178 | |
| 179 | def generate(llm_handle) -> list[dict[str, object]]: |
| 180 | outputs = ray.get(llm_handle.generate.remote(PROMPTS, SAMPLING_PARAMS)) |
| 181 | return [ |
| 182 | { |
| 183 | "token_ids": output.outputs[0].token_ids, |
| 184 | "text": output.outputs[0].text, |
| 185 | } |
| 186 | for output in outputs |
| 187 | ] |
| 188 | |
| 189 | |
| 190 | def print_generations(label: str, generations: list[dict[str, object]]) -> None: |
| 191 | print(f"\n{label}") |
| 192 | for prompt, generation in zip(PROMPTS, generations): |
| 193 | print( |
| 194 | f" {prompt!r} -> {generation['text']!r} " |
| 195 | f"(token_ids={generation['token_ids']})" |
| 196 | ) |
| 197 | |
| 198 | |
| 199 | def main() -> None: |
| 200 | ray.init() |
| 201 | train_model = None |
| 202 | pg_inference = None |
| 203 | llm = None |
| 204 | try: |
| 205 | train_model = TrainModel.remote(MODEL_NAME) |
| 206 | pg_inference = placement_group( |
| 207 | [{"GPU": 1, "CPU": 0}] * INFERENCE_TP_SIZE, |
| 208 | strategy="STRICT_PACK", |
| 209 | ) |
| 210 | ray.get(pg_inference.ready()) |
| 211 | scheduling_strategy = PlacementGroupSchedulingStrategy( |
| 212 | placement_group=pg_inference, |
| 213 | placement_group_capture_child_tasks=True, |
| 214 | placement_group_bundle_index=0, |
| 215 | ) |
| 216 | llm = launch_llm(scheduling_strategy) |
| 217 | |
| 218 | before = generate(llm) |
| 219 | print_generations("BEFORE sparse update", before) |
| 220 | |
| 221 | ray.get(llm.sleep.remote(level=0)) |
| 222 | world_size = ray.get(llm.get_world_size.remote()) + 1 |
| 223 | ray.get(train_model.init_sparse_engine.remote(world_size, llm)) |
| 224 | patched_names, num_updates = ray.get(train_model.patch_and_send.remote()) |
| 225 | ray.get(llm.wake_up.remote(tags=["scheduling"])) |
| 226 | |
| 227 | after = generate(llm) |
| 228 | print_generations("AFTER sparse update", after) |
| 229 | print(f"patched_checkpoint_names={patched_names}") |
| 230 | print(f"num_sparse_values={num_updates}") |
| 231 | outputs_changed = any( |
| 232 | old["token_ids"] != new["token_ids"] |
| 233 | for old, new in zip(before, after, strict=True) |
| 234 | ) |
| 235 | print(f"outputs_changed={outputs_changed}") |
| 236 | finally: |
| 237 | if train_model is not None: |
| 238 | with suppress(Exception): |
| 239 | ray.get(train_model.shutdown_engine.remote()) |
| 240 | if llm is not None: |
| 241 | with suppress(Exception): |
| 242 | ray.kill(llm) |
| 243 | if train_model is not None: |
| 244 | with suppress(Exception): |
| 245 | ray.kill(train_model) |
| 246 | if pg_inference is not None: |
| 247 | with suppress(Exception): |
| 248 | ray.util.remove_placement_group(pg_inference) |
| 249 | ray.shutdown() |
| 250 | |
| 251 | |
| 252 | if __name__ == "__main__": |
| 253 | main() |