feat: enable multi-process mode when running VLM model. (#330)

Author: wly-115

examples/generate_vlm.py

@@ -1,4 +1,4 @@
-# python generate_vlm.py --model /path/to/Qwen2.5-VL-7B-Instruct/ --disable_prefix_cache --disable_chunked_prefill --max_seqs_per_batch 4
+# python generate_vlm.py --model /path/to/Qwen2.5-VL-7B-Instruct/ --disable_prefix_cache --disable_chunked_prefill --max_seqs_per_batch 4 --devices='npu:0' --enable_shm
 
 import os
 import signal

xllm/core/common/options.cpp

@@ -23,7 +23,8 @@ std::string Options::to_string() const {
      << ", devices: " << devices().value_or("null")
      << ", draft_model_path: " << draft_model_path().value_or("null")
      << ", draft_devices: " << draft_devices().value_or("null")
-     << ",limit_image_per_prompt: " << limit_image_per_prompt()
+     << ", backend: " << backend()
+     << ", limit_image_per_prompt: " << limit_image_per_prompt()
      << ", block_size: " << block_size()
      << ", max_cache_size: " << max_cache_size()
      << ", max_memory_utilization: " << max_memory_utilization()

xllm/core/common/options.h

@@ -45,6 +45,9 @@ class Options {
 
   PROPERTY(std::optional<std::string>, draft_devices);
 
+  // model backend
+  PROPERTY(std::string, backend);
+
   // max image num per prompt, default 4
   PROPERTY(int32_t, limit_image_per_prompt) = 4;
 

xllm/core/distributed_runtime/dist_manager.cpp

@@ -31,65 +31,10 @@ namespace xllm {
 
 DistManager::DistManager(const runtime::Options& options) {
   auto master_node_addr = options.master_node_addr().value_or("");
-  // Single-Node Worker Mode
-  if (master_node_addr.empty()) {
-    setup_single_node_workers(options);
-  } else {
-    // Multi-node Worker Mode
+  if (!master_node_addr.empty()) {
     setup_multi_node_workers(options, master_node_addr);
-  }
-}
-
-void DistManager::setup_single_node_workers(const runtime::Options& options) {
-  const auto& devices = options.devices();
-  CHECK_EQ((devices.size() % options.dp_size()), 0)
-      << "Device size must be divisible by dp size in single-node serving "
-         "mode.";
-
-  // initialize process groups if there are multiple devices
-  if (devices.size() > 1) {
-    // create a process group for each device if there are multiple gpus
-    process_groups_ = parallel_state::create_npu_process_groups(devices);
-  }
-
-  const int32_t dp_local_tp_size = devices.size() / options.dp_size();
-  if (options.dp_size() > 1 && options.dp_size() < devices.size()) {
-    dp_local_process_groups_.reserve(options.dp_size());
-    for (size_t dp_rank = 0; dp_rank < options.dp_size(); ++dp_rank) {
-      auto dp_local_group_device_begin_idx = devices.begin();
-      std::advance(dp_local_group_device_begin_idx, dp_rank * dp_local_tp_size);
-      auto dp_local_group_device_end_idx = devices.begin();
-      std::advance(dp_local_group_device_end_idx,
-                   (dp_rank + 1) * dp_local_tp_size);
-      std::vector<torch::Device> dp_local_group_devices;
-      std::copy(dp_local_group_device_begin_idx,
-                dp_local_group_device_end_idx,
-                std::back_inserter(dp_local_group_devices));
-      dp_local_process_groups_.emplace_back(
-          parallel_state::create_npu_process_groups(dp_local_group_devices));
-    }
-  }
-
-  // create a worker(as worker client also) for each device
-  const int32_t world_size = static_cast<int32_t>(devices.size());
-  WorkerType worker_type =
-      (options.task_type() == "generate") ? WorkerType::LLM : WorkerType::ELM;
-  for (size_t i = 0; i < devices.size(); ++i) {
-    const int32_t rank = static_cast<int32_t>(i);
-    ProcessGroup* pg = world_size > 1 ? process_groups_[i].get() : nullptr;
-    // dp local process groups
-    ProcessGroup* dp_local_pg =
-        (options.dp_size() > 1 && options.dp_size() < world_size)
-            ? (dp_local_process_groups_[i / dp_local_tp_size]
-                                       [i % dp_local_tp_size])
-                  .get()
-            : nullptr;
-    ParallelArgs parallel_args(
-        rank, world_size, pg, dp_local_pg, options.dp_size());
-    workers_.emplace_back(std::make_unique<Worker>(
-        parallel_args, devices[i], options, worker_type));
-    worker_clients_.emplace_back(
-        std::make_unique<WorkerClient>(workers_.back().get()));
+  } else {
+    LOG(FATAL) << "master_node_addr is empty.";
   }
 }
 
@@ -166,10 +111,17 @@ void DistManager::setup_multi_node_workers(
 
   runtime::Options worker_server_options = options;
   worker_server_options.world_size(world_size);
-
-  WorkerType worker_type =
-      (options.task_type() == "generate") ? WorkerType::LLM : WorkerType::ELM;
-
+  WorkerType worker_type("LLM");
+  const auto& model_backend = options.backend();
+  if (model_backend == "llm") {
+    worker_type =
+        (options.task_type() == "generate") ? WorkerType::LLM : WorkerType::ELM;
+  } else if (model_backend == "vlm") {
+    worker_type = (options.task_type() == "generate") ? WorkerType::VLM
+                                                      : WorkerType::EVLM;
+  } else {
+    LOG(ERROR) << "Unsupported " << model_backend << " in multi-node.";
+  }
   // create local workers
   for (size_t i = 0; i < devices.size(); ++i) {
     // worldsize = 8

xllm/core/distributed_runtime/dist_manager.h

@@ -34,8 +34,6 @@ class DistManager {
 
  private:
   DISALLOW_COPY_AND_ASSIGN(DistManager);
-
-  void setup_single_node_workers(const runtime::Options& options);
   void setup_multi_node_workers(const runtime::Options& options,
                                 const std::string& master_node_addr);
 

xllm/core/distributed_runtime/worker_server.cpp

@@ -36,6 +36,7 @@ limitations under the License.
 #include "framework/parallel_state/collective_communicator.h"
 #include "framework/parallel_state/mapping_npu.h"
 #include "framework/state_dict/state_dict.h"
+#include "runtime/forward_params.h"
 #include "runtime/worker.h"
 #include "server/xllm_server_registry.h"
 #include "util/net.h"
@@ -65,6 +66,7 @@ void WorkerServer::create_server(
     int32_t dp_size,
     int local_rank,
     int32_t ep_size,
+    WorkerType worker_type,
     std::unique_ptr<ForwardSharedMemoryManager> input_shm_manager,
     std::unique_ptr<ForwardSharedMemoryManager> output_shm_manager) {
   Device device(d);
@@ -106,11 +108,6 @@ void WorkerServer::create_server(
   comm.create_process_groups(master_node_addr, device);
 #endif
 
-  WorkerType worker_type =
-      (options.task_type() == "generate") ? WorkerType::LLM : WorkerType::ELM;
-  CHECK(worker_type == WorkerType::LLM || worker_type == WorkerType::ELM)
-      << "Multi Node only support LLM and ELM Now, but get task type = "
-      << options.task_type();
   std::unique_ptr<Worker> worker =
       std::make_unique<Worker>(*parallel_args, device, options, worker_type);
   worker_service->set_worker(std::move(worker));
@@ -216,8 +213,8 @@ WorkerServer::WorkerServer(int local_worker_idx,
                            const runtime::Options& options,
                            WorkerType worker_type,
                            bool use_spawn_worker) {
-  if (worker_type == WorkerType::LLM || worker_type == WorkerType::ELM) {
-    // TODO: Refactor these code later.
+  if (worker_type == WorkerType::LLM || worker_type == WorkerType::ELM ||
+      worker_type == WorkerType::VLM || worker_type == WorkerType::EVLM) {
     if (use_spawn_worker) {
       // start worker in a spawn process(for offline inference worker.)
       create_spawn_server(local_worker_idx,
@@ -251,6 +248,7 @@ WorkerServer::WorkerServer(int local_worker_idx,
                                         parallel_args.dp_size(),
                                         local_worker_idx,
                                         parallel_args.ep_size(),
+                                        worker_type,
                                         std::move(input_shm_manager),
                                         std::move(output_shm_manager));
     }

xllm/core/distributed_runtime/worker_server.h

@@ -58,6 +58,7 @@ class WorkerServer {
       int32_t dp_size,
       int local_rank,
       int32_t ep_size,
+      WorkerType worker_type,
       std::unique_ptr<ForwardSharedMemoryManager> input_shm_manager,
       std::unique_ptr<ForwardSharedMemoryManager> output_shm_manager);
 

xllm/core/distributed_runtime/worker_service.cpp

@@ -12,7 +12,6 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
-
 #include "worker_service.h"
 
 #include <brpc/closure_guard.h>

xllm/core/framework/batch/batch.cpp

@@ -52,8 +52,8 @@ void Batch::add(Sequence* sequence, uint32_t allowed_max_token) {
     input_embeddings_vec_.emplace_back(input_embedding);
 
   const auto& mm_data = sequence->get_mm_data();
-  // if (sequence->is_prefill_stage() &&  mm_data.valid()) // TODO:Compatible
-  // With Chunked Prefill
+  //  if (sequence->is_prefill_stage() &&  mm_data.valid()) // TODO:Compatible
+  //  With Chunked Prefill
   if ((sequence->kv_state().kv_cache_tokens_num() <
        sequence->num_prompt_tokens()) &&
       mm_data.valid())
@@ -83,6 +83,7 @@ ForwardInput Batch::prepare_forward_input(uint32_t num_decoding_tokens,
 
 RawForwardInput Batch::prepare_forward_input(uint32_t start_idx,
                                              uint32_t end_idx,
+                                             const ModelArgs& args,
                                              ThreadPool* thread_pool) {
   BatchInputBuilder builder(sequences_,
                             allowed_max_tokens_,
@@ -91,7 +92,7 @@ RawForwardInput Batch::prepare_forward_input(uint32_t start_idx,
                             copy_in_cache_block_infos_,
                             copy_out_cache_block_infos_,
                             swap_cache_block_infos_,
-                            nullptr,
+                            &args,
                             thread_pool);
   return builder.build_raw_forward_input(start_idx, end_idx);
 }

xllm/core/framework/batch/batch.h

@@ -77,7 +77,8 @@ class Batch {
   // Convert Batch to pb type, which will be pass to remote worker.
   RawForwardInput prepare_forward_input(uint32_t start_idx,
                                         uint32_t end_idx,
-                                        ThreadPool* thread_pool = nullptr);
+                                        const ModelArgs& args,
+                                        ThreadPool* thread_pool);
 
   // process output
   //