Support Data Parallel

- enable profile run

Signed-off-by: Wuxun Zhang <[email protected]>

Author: wuxun-zhang

examples/data_parallel.py

@@ -0,0 +1,254 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Usage:
+Single node:
+    python examples/offline_inference/data_parallel.py \
+            --model="ibm-research/PowerMoE-3b" \
+            --dp-size=2 \
+            --tp-size=2
+
+Multi-node:
+    Node 0 (assume the node has ip of 10.99.48.128):
+            python examples/offline_inference/data_parallel.py \
+                    --model="ibm-research/PowerMoE-3b" \
+                    --dp-size=2 \
+                    --tp-size=2 \
+                    --node-size=2 \
+                    --node-rank=0 \
+                    --master-addr=10.99.48.128 \
+                    --master-port=13345
+    Node 1:
+            python examples/offline_inference/data_parallel.py \
+                    --model="ibm-research/PowerMoE-3b" \
+                    --dp-size=2 \
+                    --tp-size=2 \
+                    --node-size=2 \
+                    --node-rank=1 \
+                    --master-addr=10.99.48.128 \
+                    --master-port=13345
+"""
+
+import os
+from time import sleep
+import torch
+
+from vllm import LLM, SamplingParams
+from vllm.utils import get_open_port
+
+
+def parse_args():
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Data Parallel Inference")
+    parser.add_argument(
+        "--model",
+        type=str,
+        default="ibm-research/PowerMoE-3b",
+        help="Model name or path",
+    )
+    parser.add_argument(
+        "--dp-size", type=int, default=2, help="Data parallel size"
+    )
+    parser.add_argument(
+        "--tp-size", type=int, default=2, help="Tensor parallel size"
+    )
+    parser.add_argument(
+        "--node-size", type=int, default=1, help="Total number of nodes"
+    )
+    parser.add_argument(
+        "--node-rank", type=int, default=0, help="Rank of the current node"
+    )
+    parser.add_argument(
+        "--master-addr", type=str, default="", help="Master node IP address"
+    )
+    parser.add_argument(
+        "--master-port", type=int, default=0, help="Master node port"
+    )
+    parser.add_argument(
+        "--enforce-eager",
+        action="store_true",
+        help="Enforce eager mode execution.",
+    )
+    parser.add_argument(
+        "--trust-remote-code", action="store_true", help="Trust remote code."
+    )
+    parser.add_argument(
+        "--max-num-seqs",
+        type=int,
+        default=64,
+        help=(
+            "Maximum number of sequences to be processed in a single iteration."
+        ),
+    )
+    parser.add_argument(
+        "--gpu-memory-utilization",
+        type=float,
+        default=0.8,
+        help=("Fraction of GPU memory vLLM is allowed to allocate (0.0, 1.0]."),
+    )
+    parser.add_argument(
+        "--random-input",
+        action="store_true",
+        help="Use random generated input tokens.",
+    )
+    return parser.parse_args()
+
+
+def generate_random_token_ids(repeat=1) -> list[int]:
+    """
+    For testing different seuquence length in data parallel scenario
+    """
+    candidate_lens = [130, 560]
+    prompts = []
+    for num_tokens in candidate_lens:
+        tokens = torch.randint(
+            low=0, high=10000, size=(num_tokens,), dtype=torch.int32
+        )
+        [prompts.append(tokens.tolist()) for _ in range(repeat)]
+    return prompts
+
+
+def main(
+    model,
+    dp_size,
+    local_dp_rank,
+    global_dp_rank,
+    dp_master_ip,
+    dp_master_port,
+    GPUs_per_dp_rank,
+    enforce_eager,
+    trust_remote_code,
+    max_num_seqs,
+    gpu_memory_utilization,
+):
+    os.environ["VLLM_DP_RANK"] = str(global_dp_rank)
+    os.environ["VLLM_DP_RANK_LOCAL"] = str(local_dp_rank)
+    os.environ["VLLM_DP_SIZE"] = str(dp_size)
+    os.environ["VLLM_DP_MASTER_IP"] = dp_master_ip
+    os.environ["VLLM_DP_MASTER_PORT"] = str(dp_master_port)
+
+    # CUDA_VISIBLE_DEVICES for each DP rank is set automatically inside the
+    # engine processes.
+
+    # Sample prompts.
+    prompts = [
+        "Hello, my name is",
+        "The president of the United States is",
+        "The capital of France is",
+        "The future of AI is",
+    ] * 40
+
+    # generate prompts with different length to demonstrate DP aware padding.
+    if args.random_input:
+        prompts = generate_random_token_ids(40)
+
+    # with DP, each rank should process different prompts.
+    # usually all the DP ranks process a full dataset,
+    # and each rank processes a different part of the dataset.
+    floor = len(prompts) // dp_size
+    remainder = len(prompts) % dp_size
+
+    # Distribute prompts into even groups.
+    def start(rank):
+        return rank * floor + min(rank, remainder)
+
+    prompts = prompts[start(global_dp_rank) : start(global_dp_rank + 1)]
+    if len(prompts) == 0:
+        # if any rank has no prompts to process,
+        # we need to set a placeholder prompt
+        prompts = ["Placeholder"]
+    print(f"DP rank {global_dp_rank} needs to process {len(prompts)} prompts")
+    # Create a sampling params object.
+    # since we are doing data parallel, every rank can have different
+    # sampling params. here we set different max_tokens for different
+    # ranks for demonstration.
+    sampling_params = SamplingParams(
+        temperature=0.8, top_p=0.95, max_tokens=[16, 20][global_dp_rank % 2]
+    )
+
+    # Create an LLM.
+    llm = LLM(
+        model=model,
+        tensor_parallel_size=GPUs_per_dp_rank,
+        enforce_eager=enforce_eager,
+        enable_expert_parallel=True,
+        trust_remote_code=trust_remote_code,
+        max_num_seqs=max_num_seqs,
+        gpu_memory_utilization=gpu_memory_utilization,
+    )
+    if not args.random_input:
+        outputs = llm.generate(prompts, sampling_params)
+    else:
+        outputs = llm.generate(None, sampling_params, prompts)
+    # Print the outputs.
+    for i, output in enumerate(outputs):
+        if i >= 5:
+            # print only 5 outputs
+            break
+        prompt = output.prompt
+        generated_text = output.outputs[0].text
+        print(
+            f"DP rank {global_dp_rank}, Prompt: {prompt!r}, "
+            f"Generated text: {generated_text!r}"
+        )
+
+    # Give engines time to pause their processing loops before exiting.
+    sleep(1)
+
+
+if __name__ == "__main__":
+    args = parse_args()
+
+    dp_size = args.dp_size
+    tp_size = args.tp_size
+    node_size = args.node_size
+    node_rank = args.node_rank
+
+    if node_size == 1:
+        dp_master_ip = "127.0.0.1"
+        dp_master_port = get_open_port()
+    else:
+        dp_master_ip = args.master_addr
+        dp_master_port = args.master_port
+
+    assert dp_size % node_size == 0, "dp_size should be divisible by node_size"
+    dp_per_node = dp_size // node_size
+
+    from multiprocessing import Process
+
+    procs = []
+    for local_dp_rank, global_dp_rank in enumerate(
+        range(node_rank * dp_per_node, (node_rank + 1) * dp_per_node)
+    ):
+        proc = Process(
+            target=main,
+            args=(
+                args.model,
+                dp_size,
+                local_dp_rank,
+                global_dp_rank,
+                dp_master_ip,
+                dp_master_port,
+                tp_size,
+                args.enforce_eager,
+                args.trust_remote_code,
+                args.max_num_seqs,
+                args.gpu_memory_utilization,
+            ),
+        )
+        proc.start()
+        procs.append(proc)
+    exit_code = 0
+    for proc in procs:
+        proc.join(timeout=300)
+        if proc.exitcode is None:
+            print(
+                f"Killing process {proc.pid} that didn't stop within 5 minutes."
+            )
+            proc.kill()
+            exit_code = 1
+        elif proc.exitcode:
+            exit_code = proc.exitcode
+
+    exit(exit_code)

vllm_gaudi/distributed/device_communicators/hpu_communicator.py

@@ -5,10 +5,30 @@
 
 from vllm.distributed.device_communicators.base_device_communicator \
     import DeviceCommunicatorBase
+from vllm.distributed.parallel_state import get_dp_group
+from vllm.forward_context import get_forward_context
 
 import habana_frameworks.torch as htorch  # noqa: F401
 
 
+def naive_multicast(x: torch.Tensor,
+                    cu_tokens_across_dp_cpu: torch.Tensor) -> torch.Tensor:
+    assert x.dim() == 2, "Input tensor must be 2D"
+    dp_rank = get_dp_group().rank_in_group
+    dp_world_size = get_dp_group().world_size
+    buffer = torch.empty((cu_tokens_across_dp_cpu[-1], x.size(1)),
+                         device=x.device,
+                         dtype=x.dtype)
+    start = 0 if dp_rank == 0 else cu_tokens_across_dp_cpu[dp_rank - 1]
+    end = cu_tokens_across_dp_cpu[dp_rank]
+    buffer[start:end, :].copy_(x)
+    for idx in range(dp_world_size):
+        start = 0 if idx == 0 else cu_tokens_across_dp_cpu[idx - 1]
+        end = cu_tokens_across_dp_cpu[idx]
+        get_dp_group().broadcast(buffer[start:end, :], idx)
+    return buffer
+
+
 class HpuCommunicator(DeviceCommunicatorBase):
 
     def all_reduce(self, input_: torch.Tensor) -> torch.Tensor:
@@ -41,3 +61,28 @@ def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
                                                input_size[dim], ) +
                                               input_size[dim + 1:])
         return output_tensor
+
+    def dispatch(
+            self, hidden_states: torch.Tensor,
+            router_logits: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        all-gather based dispatch for HPUCommunicator.
+        """
+        cu_tokens_across_dp_cpu = get_forward_context(
+        ).dp_metadata.cu_tokens_across_dp_cpu
+        hidden_states_across_dp = naive_multicast(hidden_states,
+                                                  cu_tokens_across_dp_cpu)
+        router_logits_across_dp = naive_multicast(router_logits,
+                                                  cu_tokens_across_dp_cpu)
+        return hidden_states_across_dp, router_logits_across_dp
+
+    def combine(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        dp_rank = get_dp_group().rank_in_group
+        cu_tokens_across_dp_cpu = get_forward_context(
+        ).dp_metadata.cu_tokens_across_dp_cpu
+        start = 0 if dp_rank == 0 else cu_tokens_across_dp_cpu[dp_rank - 1]
+        end = cu_tokens_across_dp_cpu[dp_rank]
+
+        all_hidden_states = get_dp_group().all_reduce(hidden_states)
+        hidden_states = all_hidden_states[start:end, :]
+        return hidden_states

vllm_gaudi/platform.py

@@ -31,6 +31,7 @@ class HpuPlatform(Platform):
     supported_quantization: list[str] = [
         "compressed-tensors", "fp8", "inc", "awq_hpu", "gptq_hpu"
     ]
+    simple_compile_backend = "hpu_backend"
 
     @classmethod
     def get_attn_backend_cls(cls, selected_backend: _Backend, head_size: int,