wehub-resource-sync
530 lines
19 KiB
Python
530 lines
19 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""
|
|
Demonstrates dense-vs-sparse NCCL weight syncing with a real model.
|
|
|
|
This example mirrors the validation story used for the sparse NCCL MVP:
|
|
both the dense update path and the sparse patch path start from the same real
|
|
checkpoint and apply the same deterministic trainer-side patch. The script then
|
|
checks that greedy 1-token outputs match between the dense and sparse vLLM
|
|
engines after the update.
|
|
|
|
The example performs the following steps:
|
|
* Load a training model on one GPU via a Ray actor.
|
|
* Launch a vLLM engine with the same real model on a second GPU.
|
|
* Verify trainer vs vLLM baseline agreement before any update.
|
|
* Apply a deterministic patch to ``model.embed_tokens.weight`` on the trainer.
|
|
* Run a dense NCCL update into a fresh vLLM engine and collect post-update
|
|
outputs.
|
|
* Reset the trainer back to the baseline checkpoint.
|
|
* Apply the same deterministic patch again.
|
|
* Run a sparse NCCL update into another fresh vLLM engine and collect
|
|
post-update outputs.
|
|
* Compare dense vs sparse baseline outputs, dense vs sparse post-update
|
|
outputs, estimated payload sizes, and trainer-side send times.
|
|
|
|
Current sparse weight transfer MVP limitations:
|
|
* ``TP=1`` and ``PP=1`` only
|
|
* sparse updates use runtime/kernel-format parameter names
|
|
* sparse updates are not composable with checkpoint-format or packed updates
|
|
|
|
This example assumes a single-node cluster with two GPUs.
|
|
"""
|
|
|
|
import hashlib
|
|
import os
|
|
import time
|
|
from collections.abc import Sequence
|
|
|
|
import ray
|
|
import torch
|
|
from ray.util.placement_group import placement_group
|
|
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
|
|
from vllm import LLM, SamplingParams
|
|
from vllm.config import WeightTransferConfig
|
|
from vllm.distributed.weight_transfer.nccl_engine import (
|
|
NCCLTrainerSendWeightsArgs,
|
|
NCCLWeightTransferEngine,
|
|
)
|
|
from vllm.distributed.weight_transfer.sparse_nccl_engine import (
|
|
SparseNCCLWeightTransferEngine,
|
|
SparseWeightPatch,
|
|
)
|
|
from vllm.utils.network_utils import get_ip, get_open_port
|
|
|
|
MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct"
|
|
PATCHED_PARAM_NAME = "model.embed_tokens.weight"
|
|
MAX_PATCH_ROWS = 32
|
|
PROMPTS = [
|
|
"Hello, my name is",
|
|
"The president of the United States is",
|
|
"The capital of France is",
|
|
"The future of AI is",
|
|
]
|
|
SAMPLING_PARAMS = SamplingParams(temperature=0.0, max_tokens=1)
|
|
|
|
|
|
class MyLLM(LLM):
|
|
"""Configure the vLLM worker for Ray placement group execution."""
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
os.environ["VLLM_RAY_BUNDLE_INDICES"] = "0"
|
|
super().__init__(*args, **kwargs)
|
|
|
|
|
|
@ray.remote(num_gpus=1)
|
|
class TrainModel:
|
|
"""Ray actor that owns the trainer-side model and deterministic patch state."""
|
|
|
|
def __init__(self, model_name: str):
|
|
self.model_name = model_name
|
|
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
|
|
if self.tokenizer.pad_token_id is None:
|
|
self.tokenizer.pad_token = self.tokenizer.eos_token
|
|
|
|
self.model = None
|
|
self.patched_param = None
|
|
self.pending_sparse_patches: list[SparseWeightPatch] | None = None
|
|
self.model_update_group = None
|
|
self.master_address = get_ip()
|
|
self.port = get_open_port()
|
|
self.reset_model()
|
|
|
|
def reset_model(self) -> None:
|
|
self.model = AutoModelForCausalLM.from_pretrained(
|
|
self.model_name,
|
|
torch_dtype=torch.bfloat16,
|
|
).to("cuda:0")
|
|
self.model.eval()
|
|
|
|
try:
|
|
self.patched_param = self.model.get_parameter(PATCHED_PARAM_NAME)
|
|
except AttributeError as exc:
|
|
raise RuntimeError(
|
|
f"Expected trainer model to expose `{PATCHED_PARAM_NAME}`"
|
|
) from exc
|
|
|
|
self.pending_sparse_patches = None
|
|
|
|
def create_rendezvous(self) -> tuple[str, int]:
|
|
self.port = get_open_port()
|
|
return self.master_address, self.port
|
|
|
|
def init_weight_transfer_group(self, world_size: int) -> None:
|
|
self.model_update_group = NCCLWeightTransferEngine.trainer_init(
|
|
dict(
|
|
master_address=self.master_address,
|
|
master_port=self.port,
|
|
world_size=world_size,
|
|
)
|
|
)
|
|
|
|
def get_dense_update_info(self, packed: bool = False) -> tuple[dict, int]:
|
|
names = []
|
|
dtype_names = []
|
|
shapes = []
|
|
payload_bytes = 0
|
|
for name, param in self.model.named_parameters():
|
|
names.append(name)
|
|
dtype_names.append(str(param.dtype).split(".")[-1])
|
|
shapes.append(list(param.shape))
|
|
payload_bytes += param.numel() * param.element_size()
|
|
|
|
return (
|
|
dict(
|
|
names=names,
|
|
dtype_names=dtype_names,
|
|
shapes=shapes,
|
|
packed=packed,
|
|
),
|
|
payload_bytes,
|
|
)
|
|
|
|
@torch.inference_mode()
|
|
def generate(
|
|
self,
|
|
prompts: Sequence[str],
|
|
max_new_tokens: int = 1,
|
|
) -> list[dict[str, object]]:
|
|
generations = []
|
|
for prompt in prompts:
|
|
model_inputs = self.tokenizer(prompt, return_tensors="pt").to("cuda:0")
|
|
output = self.model.generate(
|
|
**model_inputs,
|
|
max_new_tokens=max_new_tokens,
|
|
do_sample=False,
|
|
pad_token_id=self.tokenizer.pad_token_id,
|
|
)
|
|
new_token_ids = output[0, model_inputs["input_ids"].shape[1] :].tolist()
|
|
generations.append(
|
|
{
|
|
"token_ids": new_token_ids,
|
|
"text": self.tokenizer.decode(
|
|
new_token_ids,
|
|
skip_special_tokens=False,
|
|
),
|
|
}
|
|
)
|
|
return generations
|
|
|
|
def prepare_sparse_patch(
|
|
self,
|
|
prompts: Sequence[str],
|
|
max_patch_rows: int = MAX_PATCH_ROWS,
|
|
) -> tuple[dict[str, object], list[int], str, int]:
|
|
selected_token_ids: list[int] = []
|
|
special_ids = set(self.tokenizer.all_special_ids)
|
|
for prompt in prompts:
|
|
token_ids = self.tokenizer(prompt, add_special_tokens=False)["input_ids"]
|
|
for token_id in token_ids:
|
|
if token_id in special_ids or token_id in selected_token_ids:
|
|
continue
|
|
selected_token_ids.append(token_id)
|
|
if len(selected_token_ids) == max_patch_rows:
|
|
break
|
|
if len(selected_token_ids) == max_patch_rows:
|
|
break
|
|
|
|
if not selected_token_ids:
|
|
raise ValueError("Could not derive any non-special token IDs to patch")
|
|
|
|
vocab_size = self.patched_param.shape[0]
|
|
next_token_id = selected_token_ids[-1]
|
|
while len(selected_token_ids) < max_patch_rows:
|
|
next_token_id = (next_token_id + 1) % vocab_size
|
|
if next_token_id in special_ids or next_token_id in selected_token_ids:
|
|
continue
|
|
selected_token_ids.append(next_token_id)
|
|
|
|
row_ids = torch.tensor(
|
|
selected_token_ids,
|
|
device=self.patched_param.device,
|
|
dtype=torch.long,
|
|
)
|
|
hidden_size = self.patched_param.shape[1]
|
|
column_offsets = torch.arange(
|
|
hidden_size,
|
|
device=self.patched_param.device,
|
|
dtype=torch.long,
|
|
)
|
|
|
|
with torch.no_grad():
|
|
# Rotate the selected embedding rows instead of zeroing them so the
|
|
# patch remains deterministic while avoiding a degenerate collapse
|
|
# to the same special token after the update.
|
|
replacement_rows = self.patched_param[row_ids].roll(shifts=1, dims=0)
|
|
self.patched_param[row_ids] = replacement_rows
|
|
|
|
flat_indices = (
|
|
row_ids.unsqueeze(1).mul(hidden_size).add(column_offsets).reshape(-1)
|
|
)
|
|
flat_values = self.patched_param[row_ids].reshape(-1).contiguous()
|
|
self.pending_sparse_patches = [
|
|
SparseWeightPatch(
|
|
name=PATCHED_PARAM_NAME,
|
|
indices=flat_indices.to(torch.int32),
|
|
values=flat_values,
|
|
)
|
|
]
|
|
patch_digest = hashlib.sha256(
|
|
self.pending_sparse_patches[0].indices.cpu().numpy().tobytes()
|
|
+ self.pending_sparse_patches[0]
|
|
.values.detach()
|
|
.float()
|
|
.cpu()
|
|
.numpy()
|
|
.tobytes()
|
|
).hexdigest()
|
|
|
|
sparse_payload_bytes = (
|
|
flat_indices.numel() * torch.tensor([], dtype=torch.int32).element_size()
|
|
+ flat_values.numel() * flat_values.element_size()
|
|
)
|
|
update_info = dict(
|
|
names=[PATCHED_PARAM_NAME],
|
|
dtype_names=[str(self.patched_param.dtype).split(".")[-1]],
|
|
shapes=[list(self.patched_param.shape)],
|
|
num_updates_list=[flat_indices.numel()],
|
|
)
|
|
return update_info, selected_token_ids, patch_digest, sparse_payload_bytes
|
|
|
|
def broadcast_weights(self, packed: bool = False) -> float:
|
|
if self.model_update_group is None:
|
|
raise RuntimeError("Weight transfer group is not initialized")
|
|
|
|
trainer_args = NCCLTrainerSendWeightsArgs(
|
|
group=self.model_update_group,
|
|
packed=packed,
|
|
)
|
|
start = time.perf_counter()
|
|
NCCLWeightTransferEngine.trainer_send_weights(
|
|
iterator=self.model.named_parameters(),
|
|
trainer_args=trainer_args,
|
|
)
|
|
torch.accelerator.synchronize()
|
|
return (time.perf_counter() - start) * 1000.0
|
|
|
|
def broadcast_pending_sparse_patch(self) -> float:
|
|
if self.model_update_group is None:
|
|
raise RuntimeError("Weight transfer group is not initialized")
|
|
if self.pending_sparse_patches is None:
|
|
raise RuntimeError("Sparse patch has not been prepared")
|
|
|
|
start = time.perf_counter()
|
|
SparseNCCLWeightTransferEngine.trainer_send_weights(
|
|
iter(self.pending_sparse_patches),
|
|
NCCLTrainerSendWeightsArgs(group=self.model_update_group),
|
|
)
|
|
torch.accelerator.synchronize()
|
|
self.pending_sparse_patches = None
|
|
return (time.perf_counter() - start) * 1000.0
|
|
|
|
|
|
def launch_llm(
|
|
scheduling_inference: PlacementGroupSchedulingStrategy,
|
|
backend: str = "nccl",
|
|
):
|
|
return ray.remote(
|
|
num_cpus=0,
|
|
num_gpus=0,
|
|
scheduling_strategy=scheduling_inference,
|
|
)(MyLLM).remote(
|
|
model=MODEL_NAME,
|
|
enforce_eager=True,
|
|
tensor_parallel_size=1,
|
|
distributed_executor_backend="ray",
|
|
gpu_memory_utilization=0.7,
|
|
weight_transfer_config=WeightTransferConfig(backend=backend),
|
|
)
|
|
|
|
|
|
def collect_vllm_generations(llm_handle) -> list[dict[str, object]]:
|
|
outputs = ray.get(llm_handle.generate.remote(PROMPTS, SAMPLING_PARAMS))
|
|
generations = []
|
|
for output in outputs:
|
|
generations.append(
|
|
{
|
|
"token_ids": output.outputs[0].token_ids,
|
|
"text": output.outputs[0].text,
|
|
}
|
|
)
|
|
return generations
|
|
|
|
|
|
def token_sequences_match(
|
|
left: Sequence[dict[str, object]],
|
|
right: Sequence[dict[str, object]],
|
|
) -> bool:
|
|
return [item["token_ids"] for item in left] == [item["token_ids"] for item in right]
|
|
|
|
|
|
def print_generations(label: str, prompts: Sequence[str], generations) -> None:
|
|
print(f"\n{label}")
|
|
print("-" * 50)
|
|
for prompt, generation in zip(prompts, generations):
|
|
print(f"Prompt: {prompt!r}")
|
|
print(f"Token IDs: {generation['token_ids']}")
|
|
print(f"Text: {generation['text']!r}")
|
|
print("-" * 50)
|
|
|
|
|
|
def run_dense_phase(
|
|
train_model,
|
|
scheduling_inference: PlacementGroupSchedulingStrategy,
|
|
) -> dict[str, object]:
|
|
ray.get(train_model.reset_model.remote())
|
|
llm = launch_llm(scheduling_inference, backend="nccl")
|
|
try:
|
|
dense_before = collect_vllm_generations(llm)
|
|
|
|
ray.get(llm.sleep.remote(level=0))
|
|
master_address, master_port = ray.get(train_model.create_rendezvous.remote())
|
|
world_size = ray.get(llm.get_world_size.remote()) + 1
|
|
inference_init = llm.init_weight_transfer_engine.remote(
|
|
dict(
|
|
init_info=dict(
|
|
master_address=master_address,
|
|
master_port=master_port,
|
|
rank_offset=1,
|
|
world_size=world_size,
|
|
)
|
|
)
|
|
)
|
|
trainer_init = train_model.init_weight_transfer_group.remote(world_size)
|
|
ray.get([trainer_init, inference_init])
|
|
ray.get(llm.start_weight_update.remote())
|
|
|
|
dense_update_info, dense_payload_bytes = ray.get(
|
|
train_model.get_dense_update_info.remote()
|
|
)
|
|
_, selected_token_ids, patch_digest, _ = ray.get(
|
|
train_model.prepare_sparse_patch.remote(PROMPTS)
|
|
)
|
|
|
|
inference_update = llm.update_weights.remote(
|
|
dict(update_info=dense_update_info)
|
|
)
|
|
dense_send_ms, _ = ray.get(
|
|
[
|
|
train_model.broadcast_weights.remote(packed=False),
|
|
inference_update,
|
|
]
|
|
)
|
|
ray.get(llm.finish_weight_update.remote())
|
|
ray.get(llm.wake_up.remote(tags=["scheduling"]))
|
|
|
|
dense_after = collect_vllm_generations(llm)
|
|
|
|
return {
|
|
"dense_before": dense_before,
|
|
"dense_after": dense_after,
|
|
"selected_token_ids": selected_token_ids,
|
|
"patch_digest": patch_digest,
|
|
"dense_payload_bytes": dense_payload_bytes,
|
|
"dense_send_ms": dense_send_ms,
|
|
}
|
|
finally:
|
|
ray.kill(llm)
|
|
|
|
|
|
def run_sparse_phase(
|
|
train_model,
|
|
scheduling_inference: PlacementGroupSchedulingStrategy,
|
|
) -> dict[str, object]:
|
|
ray.get(train_model.reset_model.remote())
|
|
llm = launch_llm(scheduling_inference, backend="sparse_nccl")
|
|
try:
|
|
sparse_before = collect_vllm_generations(llm)
|
|
|
|
ray.get(llm.sleep.remote(level=0))
|
|
master_address, master_port = ray.get(train_model.create_rendezvous.remote())
|
|
world_size = ray.get(llm.get_world_size.remote()) + 1
|
|
inference_init = llm.init_weight_transfer_engine.remote(
|
|
dict(
|
|
init_info=dict(
|
|
master_address=master_address,
|
|
master_port=master_port,
|
|
rank_offset=1,
|
|
world_size=world_size,
|
|
)
|
|
)
|
|
)
|
|
trainer_init = train_model.init_weight_transfer_group.remote(world_size)
|
|
ray.get([trainer_init, inference_init])
|
|
ray.get(llm.start_weight_update.remote())
|
|
|
|
sparse_update_info, selected_token_ids, patch_digest, sparse_payload_bytes = (
|
|
ray.get(train_model.prepare_sparse_patch.remote(PROMPTS))
|
|
)
|
|
|
|
inference_update = llm.update_weights.remote(
|
|
dict(update_info=sparse_update_info)
|
|
)
|
|
sparse_send_ms, _ = ray.get(
|
|
[
|
|
train_model.broadcast_pending_sparse_patch.remote(),
|
|
inference_update,
|
|
]
|
|
)
|
|
ray.get(llm.finish_weight_update.remote())
|
|
ray.get(llm.wake_up.remote(tags=["scheduling"]))
|
|
|
|
sparse_after = collect_vllm_generations(llm)
|
|
|
|
return {
|
|
"sparse_before": sparse_before,
|
|
"sparse_after": sparse_after,
|
|
"selected_token_ids": selected_token_ids,
|
|
"patch_digest": patch_digest,
|
|
"sparse_payload_bytes": sparse_payload_bytes,
|
|
"sparse_send_ms": sparse_send_ms,
|
|
}
|
|
finally:
|
|
ray.kill(llm)
|
|
|
|
|
|
ray.init()
|
|
|
|
try:
|
|
train_model = TrainModel.remote(MODEL_NAME)
|
|
|
|
pg_inference = placement_group([{"GPU": 1, "CPU": 0}])
|
|
ray.get(pg_inference.ready())
|
|
scheduling_inference = PlacementGroupSchedulingStrategy(
|
|
placement_group=pg_inference,
|
|
placement_group_capture_child_tasks=True,
|
|
placement_group_bundle_index=0,
|
|
)
|
|
|
|
dense_results = run_dense_phase(train_model, scheduling_inference)
|
|
sparse_results = run_sparse_phase(train_model, scheduling_inference)
|
|
|
|
baseline_equal = token_sequences_match(
|
|
dense_results["dense_before"],
|
|
sparse_results["sparse_before"],
|
|
)
|
|
patch_selection_equal = (
|
|
dense_results["selected_token_ids"] == sparse_results["selected_token_ids"]
|
|
)
|
|
patch_digest_equal = dense_results["patch_digest"] == sparse_results["patch_digest"]
|
|
after_equal = token_sequences_match(
|
|
dense_results["dense_after"],
|
|
sparse_results["sparse_after"],
|
|
)
|
|
any_output_changed = any(
|
|
before["token_ids"] != after["token_ids"]
|
|
for before, after in zip(
|
|
dense_results["dense_before"],
|
|
dense_results["dense_after"],
|
|
)
|
|
)
|
|
dense_payload_mb = dense_results["dense_payload_bytes"] / (1024 * 1024)
|
|
sparse_payload_mb = sparse_results["sparse_payload_bytes"] / (1024 * 1024)
|
|
|
|
print_generations(
|
|
"Dense baseline outputs",
|
|
PROMPTS,
|
|
dense_results["dense_before"],
|
|
)
|
|
print_generations(
|
|
"Sparse baseline outputs", PROMPTS, sparse_results["sparse_before"]
|
|
)
|
|
print_generations(
|
|
"Dense outputs after update", PROMPTS, dense_results["dense_after"]
|
|
)
|
|
print_generations(
|
|
"Sparse outputs after update",
|
|
PROMPTS,
|
|
sparse_results["sparse_after"],
|
|
)
|
|
|
|
print(f"patched_token_ids = {dense_results['selected_token_ids']}")
|
|
print(f"patch_selection_equal = {patch_selection_equal}")
|
|
print(f"dense_patch_digest = {dense_results['patch_digest']}")
|
|
print(f"sparse_patch_digest = {sparse_results['patch_digest']}")
|
|
print(f"patch_digest_equal = {patch_digest_equal}")
|
|
print(f"baseline_equal = {baseline_equal}")
|
|
print(f"after_equal = {after_equal}")
|
|
print(f"any_output_changed = {any_output_changed}")
|
|
print(f"dense_payload_mb = {dense_payload_mb:.2f}")
|
|
print(f"sparse_payload_mb = {sparse_payload_mb:.2f}")
|
|
print(f"dense_send_ms = {dense_results['dense_send_ms']:.2f}")
|
|
print(f"sparse_send_ms = {sparse_results['sparse_send_ms']:.2f}")
|
|
|
|
if not baseline_equal:
|
|
raise RuntimeError(
|
|
"Dense and sparse phases did not start from the same baseline"
|
|
)
|
|
if not patch_selection_equal:
|
|
raise RuntimeError("Dense and sparse phases used different sparse patches")
|
|
if not patch_digest_equal:
|
|
raise RuntimeError("Dense and sparse phases produced different patch values")
|
|
if not after_equal:
|
|
raise RuntimeError("Dense and sparse updates produced different outputs")
|
|
if not any_output_changed:
|
|
raise RuntimeError("Patch did not change the observed outputs")
|
|
finally:
|
|
ray.shutdown()
|