| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Sharded-RDT weight sync on FOUR GPUs, one node: 2 FSDP2 trainer ranks -> 2 |
| 4 | vLLM DP ranks with expert parallelism. |
| 5 | |
| 6 | The reference example for the backend, sized to fit a CI runner: |
| 7 | ``CheckpointNameSource`` over an FSDP2 model, one serve actor per rank, the |
| 8 | HTTP control plane, and assertions so it can run unattended. |
| 9 | |
| 10 | Why a MoE model rather than a small dense one: the point of sharded RDT is that |
| 11 | each inference worker pulls only the slices it consumes, so the example is only |
| 12 | meaningful when no rank holds the whole model. That needs experts on both sides. |
| 13 | ``Qwen/Qwen1.5-MoE-A2.7B`` is the reference shape, a real MoE of the same family |
| 14 | as the 30B model the big example uses. Its WEIGHTS are never read: the trainer |
| 15 | builds the config and the server runs ``--load-format dummy``, so the sync is |
| 16 | what is under test and nothing is downloaded but the config and tokenizer. A |
| 17 | smaller MoE via ``RDT_MODEL`` changes nothing about what is exercised, which is |
| 18 | what CI runs. |
| 19 | |
| 20 | The trainer publishes CHECKPOINT names, not the names its own modules carry: |
| 21 | Transformers fuses each layer's experts into ``[E, ...]`` tensors |
| 22 | (``mlp.experts.gate_up_proj``/``down_proj``) while the checkpoint stores them per |
| 23 | expert, and vLLM's MoE loaders read the per-expert form. ``CheckpointNameSource`` |
| 24 | splits them back, which is the same conversion a real trainer does when it maps |
| 25 | its internal layout onto checkpoint names. Publishing the fused names instead |
| 26 | works only for the Qwen families, whose loaders happen to accept them. |
| 27 | |
| 28 | FSDP2 shards those fused tensors on dim 0 — the expert dimension — so after |
| 29 | sharding no rank stores every expert, and the inference side gets real EP from |
| 30 | ``--enable-expert-parallel`` over DP2. Two simplifications an example can afford |
| 31 | and a real trainer cannot: each rank BUILDS the whole model before |
| 32 | ``fully_shard`` splits it (fine for a tiny config, so nothing here streams a |
| 33 | checkpoint), and iteration all-gathers every parameter, so ownership stays |
| 34 | uniform. Serving only the experts a rank holds (``held_names``) is where a real |
| 35 | trainer earns its keep and is deliberately out of scope. |
| 36 | |
| 37 | Model naming is load-bearing: names are mapped onto checkpoint keys one-to-one, |
| 38 | so a model whose HF module names have drifted from its checkpoint (as |
| 39 | GraniteMoe's router has) will not work here. |
| 40 | |
| 41 | What it checks (a smoke test, not a numerical equivalence check): |
| 42 | - generation CHANGES after the first sync — the weights really moved, so a |
| 43 | silently-skipped sync fails instead of printing plausible text; |
| 44 | - generation is UNCHANGED by a second sync — replay is stable. |
| 45 | |
| 46 | Run: |
| 47 | python examples/rl/rlhf_sharded_rdt_small_ep.py |
| 48 | |
| 49 | Needs 4 GPUs on one node, each big enough for the whole model, since every rank |
| 50 | builds it before ``fully_shard`` splits it: the default is ~29 GiB in bf16, so |
| 51 | 40 GiB cards and up. Set ``RDT_MODEL`` to a tiny MoE on smaller GPUs. Joins an |
| 52 | existing Ray cluster if there is one and otherwise starts its own. |
| 53 | """ |
| 54 | |
| 55 | import os |
| 56 | import sys |
| 57 | import time |
| 58 | |
| 59 | import ray |
| 60 | import torch |
| 61 | import torch.distributed as dist |
| 62 | from ray.util.placement_group import placement_group, placement_group_table |
| 63 | from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy |
| 64 | from rdt_vllm_serve import ( |
| 65 | http_generate, |
| 66 | launch_vllm_serve, |
| 67 | pause_generation, |
| 68 | resume_generation, |
| 69 | shutdown_server, |
| 70 | wait_for_server, |
| 71 | ) |
| 72 | from rdt_weight_source import CheckpointNameSource |
| 73 | from torch.distributed.fsdp import fully_shard |
| 74 | from transformers import AutoConfig, AutoModelForCausalLM |
| 75 | |
| 76 | from vllm.distributed.weight_transfer import ( |
| 77 | HTTPVLLMWeightSyncClient, |
| 78 | WeightTransferTrainerFactory, |
| 79 | ) |
| 80 | from vllm.distributed.weight_transfer.sharded_rdt_trainer import ( |
| 81 | ShardedRDTTrainerInitInfo, |
| 82 | ) |
| 83 | from vllm.utils.network_utils import get_open_port |
| 84 | |
| 85 | # Weights are built, never loaded, so this costs a config fetch -- but every rank |
| 86 | # materializes the model before sharding, so the default wants a 40 GiB card. |
| 87 | MODEL_NAME = os.environ.get("RDT_MODEL", "Qwen/Qwen1.5-MoE-A2.7B") |
| 88 | RAY_NAMESPACE = "sharded_rdt_small_ep_example" |
| 89 | VLLM_PORT = int(os.environ.get("RDT_VLLM_PORT", "8100")) |
| 90 | VLLM_ENDPOINT = f"http://127.0.0.1:{VLLM_PORT}" |
| 91 | |
| 92 | FSDP_WORLD_SIZE = 2 |
| 93 | INFERENCE_TP_SIZE = 1 |
| 94 | INFERENCE_DP_SIZE = 2 |
| 95 | NUM_INFERENCE_CONSUMERS = INFERENCE_TP_SIZE * INFERENCE_DP_SIZE |
| 96 | |
| 97 | PROMPTS = ["The capital of France is", "The future of AI is"] |
| 98 | |
| 99 | |
| 100 | @ray.remote(num_gpus=1) |
| 101 | class FSDPTrainWorker: |
| 102 | """One FSDP2 rank per GPU. Identical to the 8-GPU example's worker: the |
| 103 | sharded-RDT engine owns the NIXL serve surface, so this stays a plain Ray |
| 104 | actor with no producer mixin.""" |
| 105 | |
| 106 | def __init__(self, model_name, rank, world_size, master_addr, master_port): |
| 107 | self.rank = rank |
| 108 | self.engine = None |
| 109 | os.environ["MASTER_ADDR"] = master_addr |
| 110 | os.environ["MASTER_PORT"] = str(master_port) |
| 111 | dist.init_process_group(backend="nccl", rank=rank, world_size=world_size) |
| 112 | torch.accelerator.set_device_index(0) |
| 113 | |
| 114 | # Same seed on every rank, so the shards compose into ONE model rather |
| 115 | # than a blend of independent inits -- the second sync's stability check |
| 116 | # is only meaningful against coherent weights. |
| 117 | torch.manual_seed(0) |
| 118 | config = AutoConfig.from_pretrained(model_name) |
| 119 | with torch.device("cuda"): |
| 120 | model = AutoModelForCausalLM.from_config(config, dtype=torch.bfloat16) |
| 121 | for layer in model.model.layers: |
| 122 | fully_shard(layer) |
| 123 | fully_shard(model) |
| 124 | self.model = model |
| 125 | |
| 126 | def get_rank(self): |
| 127 | return self.rank |
| 128 | |
| 129 | def setup_engine(self, vllm_endpoint: str): |
| 130 | self.engine = WeightTransferTrainerFactory.trainer_init( |
| 131 | ShardedRDTTrainerInitInfo( |
| 132 | rank=self.rank, |
| 133 | num_consumers=NUM_INFERENCE_CONSUMERS, |
| 134 | trainer_actor_namespace=RAY_NAMESPACE, |
| 135 | num_rdt_buffers=int(os.environ.get("NUM_RDT_BUFFERS", "2")), |
| 136 | buffer_presize_gb=float(os.environ.get("RDT_BUFFER_PRESIZE_GB", "0")), |
| 137 | ), |
| 138 | client=HTTPVLLMWeightSyncClient(vllm_endpoint), |
| 139 | source=CheckpointNameSource(self.model), |
| 140 | ) |
| 141 | |
| 142 | def sync_weights(self): |
| 143 | self.engine.send_weights() |
| 144 | |
| 145 | |
| 146 | def main(): |
| 147 | # Ship a minimal working_dir (this example dir) so Ray actors do NOT inherit |
| 148 | # a workspace snapshot that shadows the editable vLLM install. |
| 149 | runtime_env: dict[str, object] = { |
| 150 | "py_executable": sys.executable, |
| 151 | "working_dir": os.path.dirname(os.path.abspath(__file__)), |
| 152 | } |
| 153 | forwarded = { |
| 154 | k: os.environ[k] |
| 155 | for k in ( |
| 156 | "NCCL_CUMEM_ENABLE", |
| 157 | "VLLM_NCCL_SO_PATH", |
| 158 | "LD_PRELOAD", |
| 159 | "LD_LIBRARY_PATH", |
| 160 | ) |
| 161 | if k in os.environ |
| 162 | } |
| 163 | if forwarded: |
| 164 | runtime_env["env_vars"] = forwarded |
| 165 | if not ray.is_initialized(): |
| 166 | try: |
| 167 | ray.init(address="auto", runtime_env=runtime_env, namespace=RAY_NAMESPACE) |
| 168 | except (ConnectionError, RuntimeError): |
| 169 | # No cluster to join (a bare CI runner): start one on this node. |
| 170 | # `launch_vllm_serve` hands the child this cluster's address, which |
| 171 | # it needs -- vLLM's ray DP backend would otherwise start its own. |
| 172 | ray.init(runtime_env=runtime_env, namespace=RAY_NAMESPACE) |
| 173 | |
| 174 | print( |
| 175 | f"[init] model {MODEL_NAME}, {FSDP_WORLD_SIZE} trainer + " |
| 176 | f"{INFERENCE_DP_SIZE} inference GPUs" |
| 177 | ) |
| 178 | |
| 179 | # Reserve the trainer's GPUs BEFORE launching the server, so the server's DP |
| 180 | # workers take the node's remaining two. STRICT_PACK keeps the trainer's |
| 181 | # all-gathers on one node's NVLink -- here that is the only node anyway. |
| 182 | trainer_pg = placement_group( |
| 183 | [{"GPU": 1, "CPU": 1}] * FSDP_WORLD_SIZE, strategy="STRICT_PACK" |
| 184 | ) |
| 185 | ray.get(trainer_pg.ready()) |
| 186 | pg_node_id = next( |
| 187 | iter(placement_group_table(trainer_pg)["bundles_to_node_id"].values()) |
| 188 | ) |
| 189 | master_addr = next( |
| 190 | n["NodeManagerAddress"] for n in ray.nodes() if n["NodeID"] == pg_node_id |
| 191 | ) |
| 192 | |
| 193 | @ray.remote( |
| 194 | num_cpus=0, |
| 195 | scheduling_strategy=PlacementGroupSchedulingStrategy( |
| 196 | placement_group=trainer_pg |
| 197 | ), |
| 198 | ) |
| 199 | def _free_port(): |
| 200 | return get_open_port() |
| 201 | |
| 202 | master_port = ray.get(_free_port.remote()) |
| 203 | |
| 204 | workers = [ |
| 205 | FSDPTrainWorker.options( |
| 206 | scheduling_strategy=PlacementGroupSchedulingStrategy( |
| 207 | placement_group=trainer_pg, placement_group_bundle_index=rank |
| 208 | ), |
| 209 | ).remote(MODEL_NAME, rank, FSDP_WORLD_SIZE, master_addr, master_port) |
| 210 | for rank in range(FSDP_WORLD_SIZE) |
| 211 | ] |
| 212 | ray.get([w.get_rank.remote() for w in workers]) |
| 213 | print(f"[init] {FSDP_WORLD_SIZE} FSDP trainer workers ready (seeded weights).") |
| 214 | |
| 215 | server = launch_vllm_serve( |
| 216 | MODEL_NAME, |
| 217 | tensor_parallel_size=INFERENCE_TP_SIZE, |
| 218 | data_parallel_size=INFERENCE_DP_SIZE, |
| 219 | enable_expert_parallel=True, |
| 220 | port=VLLM_PORT, |
| 221 | extra_args=["--max-model-len", "2048", "--max-num-seqs", "8"], |
| 222 | ) |
| 223 | try: |
| 224 | wait_for_server(VLLM_ENDPOINT, server) |
| 225 | |
| 226 | before = http_generate(VLLM_ENDPOINT, MODEL_NAME, PROMPTS) |
| 227 | print("[generate] BEFORE sync (vLLM's own dummy init):") |
| 228 | for p, t in zip(PROMPTS, before): |
| 229 | print(f" {p!r} -> {t!r}") |
| 230 | |
| 231 | print("[sync] building trainer engines (bake on the sender)...") |
| 232 | t0 = time.perf_counter() |
| 233 | ray.get([w.setup_engine.remote(VLLM_ENDPOINT) for w in workers]) |
| 234 | print(f"[sync] engine setup (incl. bake) took {time.perf_counter() - t0:.2f}s") |
| 235 | |
| 236 | pause_generation(VLLM_ENDPOINT) |
| 237 | t0 = time.perf_counter() |
| 238 | ray.get([w.sync_weights.remote() for w in workers]) |
| 239 | print(f"[sync] iter 0 took {time.perf_counter() - t0:.3f}s") |
| 240 | resume_generation(VLLM_ENDPOINT) |
| 241 | |
| 242 | after = http_generate(VLLM_ENDPOINT, MODEL_NAME, PROMPTS) |
| 243 | print("[generate] AFTER sync (the trainer's weights):") |
| 244 | for p, t in zip(PROMPTS, after): |
| 245 | print(f" {p!r} -> {t!r}") |
| 246 | |
| 247 | # Sync again; replay must be stable, so generation must not move. |
| 248 | pause_generation(VLLM_ENDPOINT) |
| 249 | t0 = time.perf_counter() |
| 250 | ray.get([w.sync_weights.remote() for w in workers]) |
| 251 | print(f"[sync] iter 1 took {time.perf_counter() - t0:.3f}s") |
| 252 | resume_generation(VLLM_ENDPOINT) |
| 253 | again = http_generate(VLLM_ENDPOINT, MODEL_NAME, PROMPTS) |
| 254 | |
| 255 | # Greedy decoding, so both comparisons are exact. |
| 256 | assert after != before, ( |
| 257 | f"generation did not change across the weight sync, so the sync " |
| 258 | f"moved nothing: {before!r}" |
| 259 | ) |
| 260 | assert again == after, ( |
| 261 | f"a second sync changed generation, so replay is not stable: " |
| 262 | f"{after!r} then {again!r}" |
| 263 | ) |
| 264 | print("[ok] weights moved on sync 0 and replay was stable on sync 1") |
| 265 | finally: |
| 266 | shutdown_server(server) |
| 267 | |
| 268 | |
| 269 | if __name__ == "__main__": |
| 270 | main() |