From a756b4f9ff8bf024c4fae9266a893a9f4d8fa8fe Mon Sep 17 00:00:00 2001 From: huangweizhe1 Date: Wed, 17 Jun 2026 17:42:45 +0800 Subject: [PATCH 1/3] feat: support multi machine for offline inference. --- xllm/__init__.py | 6 ++ .../core/distributed_runtime/comm_channel.cpp | 14 +++++ xllm/core/distributed_runtime/comm_channel.h | 2 + .../core/distributed_runtime/dist_manager.cpp | 44 +++++++++++++++ xllm/core/distributed_runtime/dist_manager.h | 5 ++ xllm/core/distributed_runtime/engine.h | 2 + xllm/core/distributed_runtime/llm_engine.cpp | 7 +++ xllm/core/distributed_runtime/llm_engine.h | 2 + xllm/core/distributed_runtime/llm_master.cpp | 6 ++ xllm/core/distributed_runtime/llm_master.h | 3 + xllm/core/distributed_runtime/master.h | 2 + .../distributed_runtime/remote_worker.cpp | 9 +++ xllm/core/distributed_runtime/remote_worker.h | 2 + .../speculative_engine.cpp | 7 +++ .../distributed_runtime/speculative_engine.h | 2 + xllm/core/distributed_runtime/vlm_engine.cpp | 7 +++ xllm/core/distributed_runtime/vlm_engine.h | 2 + xllm/core/distributed_runtime/vlm_master.cpp | 9 +++ xllm/core/distributed_runtime/vlm_master.h | 3 + .../distributed_runtime/worker_service.cpp | 18 ++++++ .../core/distributed_runtime/worker_service.h | 5 ++ xllm/core/runtime/worker_client.cpp | 4 ++ xllm/core/runtime/worker_client.h | 2 + xllm/proto/worker.proto | 1 + xllm/pybind/bind.cpp | 34 ++++++++++++ xllm/pybind/embedding.py | 21 ++++++- xllm/pybind/llm.py | 18 +++++- xllm/pybind/utils.py | 55 ++++++++++++++++++- xllm/pybind/vlm.py | 21 ++++++- 29 files changed, 304 insertions(+), 9 deletions(-) diff --git a/xllm/__init__.py b/xllm/__init__.py index 573abde8fe..5dce1746e6 100644 --- a/xllm/__init__.py +++ b/xllm/__init__.py @@ -61,8 +61,10 @@ def _load_xllm_export() -> ModuleType: "ArgumentParser", "Embedding", "LLM", + "LLMAssistantMaster", "LLMMaster", "VLM", + "VLMAssistantMaster", "VLMMaster", "Options", "SamplingParams", @@ -104,8 +106,10 @@ def _load_public_api() -> None: "ArgumentParser": ArgumentParser, "Embedding": Embedding, "LLM": LLM, + "LLMAssistantMaster": xllm_export.LLMAssistantMaster, "LLMMaster": xllm_export.LLMMaster, "VLM": VLM, + "VLMAssistantMaster": xllm_export.VLMAssistantMaster, "VLMMaster": xllm_export.VLMMaster, "Options": xllm_export.Options, "SamplingParams": SamplingParams, @@ -139,8 +143,10 @@ def __dir__() -> list[str]: "ArgumentParser", "Embedding", "LLM", + "LLMAssistantMaster", "LLMMaster", "VLM", + "VLMAssistantMaster", "VLMMaster", "Options", "SamplingParams", diff --git a/xllm/core/distributed_runtime/comm_channel.cpp b/xllm/core/distributed_runtime/comm_channel.cpp index 7f63f59027..d40ddacb35 100644 --- a/xllm/core/distributed_runtime/comm_channel.cpp +++ b/xllm/core/distributed_runtime/comm_channel.cpp @@ -70,6 +70,20 @@ bool CommChannel::check_health() { return resp.ok(); } +bool CommChannel::shutdown() { + proto::Empty req; + proto::Status resp; + brpc::Controller cntl; + + cntl.set_timeout_ms(5000); + stub_->Shutdown(&cntl, &req, &resp, nullptr); + if (cntl.Failed() || !resp.ok()) { + LOG(ERROR) << "Shutdown failed: " << cntl.ErrorText(); + return false; + } + return true; +} + bool CommChannel::allocate_kv_cache(const KVCacheShape& kv_cache_shape) { if (!kv_cache_shape.has_key_cache_shape() || kv_cache_shape.key_cache_shape().empty()) { diff --git a/xllm/core/distributed_runtime/comm_channel.h b/xllm/core/distributed_runtime/comm_channel.h index a0d2575be7..939ed8c1da 100644 --- a/xllm/core/distributed_runtime/comm_channel.h +++ b/xllm/core/distributed_runtime/comm_channel.h @@ -122,6 +122,8 @@ class CommChannel { virtual bool stop_profile(); + virtual bool shutdown(); + protected: bool execute_model_with_brpc( const ForwardInput& input, diff --git a/xllm/core/distributed_runtime/dist_manager.cpp b/xllm/core/distributed_runtime/dist_manager.cpp index eb66ba45d5..96e1847d1d 100644 --- a/xllm/core/distributed_runtime/dist_manager.cpp +++ b/xllm/core/distributed_runtime/dist_manager.cpp @@ -17,6 +17,8 @@ limitations under the License. #include +#include + #include "comm_channel.h" #include "common/health_check_manager.h" #include "core/framework/config/service_config.h" @@ -47,6 +49,8 @@ DistManager::DistManager(const runtime::Options& options) } DistManager::~DistManager() { + shutdown_remote_workers(); + // Stop health check HealthCheckManager::instance().stop_health_check_thread(); @@ -177,6 +181,8 @@ void DistManager::setup_multi_node_workers( const int32_t each_node_ranks = static_cast(devices.size()); const int32_t world_size = each_node_ranks * options.nnodes(); const int32_t base_rank = options.node_rank() * each_node_ranks; + local_rank_start_ = base_rank; + local_rank_count_ = each_node_ranks; const int32_t dp_size = options.dp_size(); const int32_t cp_size = options.cp_size(); const int32_t ep_size = options.ep_size(); @@ -335,4 +341,42 @@ void DistManager::setup_multi_node_workers( } } } + +bool DistManager::shutdown_remote_workers() { + if (remote_workers_shutdown_) { + return true; + } + remote_workers_shutdown_ = true; + + if (worker_clients_.empty()) { + return true; + } + + std::vector> futures; + futures.reserve(worker_clients_.size()); + const int32_t local_rank_end = local_rank_start_ + local_rank_count_; + for (size_t worker_rank = 0; worker_rank < worker_clients_.size(); + ++worker_rank) { + const int32_t global_rank = static_cast(worker_rank); + if (global_rank >= local_rank_start_ && global_rank < local_rank_end) { + continue; + } + if (local_rank_count_ > 1 && global_rank % local_rank_count_ != 0) { + continue; + } + futures.emplace_back(worker_clients_[worker_rank]->shutdown_async()); + } + if (futures.empty()) { + return true; + } + + bool success = true; + auto results = folly::collectAll(futures).get(); + for (const auto& result : results) { + if (!result.hasValue() || !result.value()) { + success = false; + } + } + return success; +} } // namespace xllm diff --git a/xllm/core/distributed_runtime/dist_manager.h b/xllm/core/distributed_runtime/dist_manager.h index f8474bfd4b..a8dd057bf3 100755 --- a/xllm/core/distributed_runtime/dist_manager.h +++ b/xllm/core/distributed_runtime/dist_manager.h @@ -32,6 +32,8 @@ class DistManager { return worker_clients_; }; + bool shutdown_remote_workers(); + private: DISALLOW_COPY_AND_ASSIGN(DistManager); void setup_multi_node_workers(const runtime::Options& options, @@ -57,5 +59,8 @@ class DistManager { std::vector> servers_; std::string server_name_; + int32_t local_rank_start_ = 0; + int32_t local_rank_count_ = 0; + bool remote_workers_shutdown_ = false; }; } // namespace xllm diff --git a/xllm/core/distributed_runtime/engine.h b/xllm/core/distributed_runtime/engine.h index ee7419b168..efbfa4f716 100644 --- a/xllm/core/distributed_runtime/engine.h +++ b/xllm/core/distributed_runtime/engine.h @@ -171,6 +171,8 @@ class Engine { return false; }; + virtual bool shutdown_remote_workers() { return true; }; + // XTensor mode: get GlobalXTensor offsets for allocated blocks // Returns per-layer K/V offsets for each block // Output: offsets[layer_id] = {k_offsets, v_offsets} diff --git a/xllm/core/distributed_runtime/llm_engine.cpp b/xllm/core/distributed_runtime/llm_engine.cpp index afe50b826d..1d5c79d7a8 100644 --- a/xllm/core/distributed_runtime/llm_engine.cpp +++ b/xllm/core/distributed_runtime/llm_engine.cpp @@ -910,6 +910,13 @@ bool LLMEngine::unlink_p2p(const std::vector& remote_addrs) { return true; } +bool LLMEngine::shutdown_remote_workers() { + if (!dist_manager_) { + return true; + } + return dist_manager_->shutdown_remote_workers(); +} + ForwardOutput LLMEngine::step(std::vector& batch) { if (worker_clients_.empty()) { // empty worker, return diff --git a/xllm/core/distributed_runtime/llm_engine.h b/xllm/core/distributed_runtime/llm_engine.h index ae2300f017..8bba5e9486 100644 --- a/xllm/core/distributed_runtime/llm_engine.h +++ b/xllm/core/distributed_runtime/llm_engine.h @@ -117,6 +117,8 @@ class LLMEngine : public Engine { bool unlink_p2p(const std::vector& remote_addrs) override; + bool shutdown_remote_workers() override; + std::shared_ptr get_dist_manager() { return dist_manager_; }; bool sleep(MasterStatus master_status) override; diff --git a/xllm/core/distributed_runtime/llm_master.cpp b/xllm/core/distributed_runtime/llm_master.cpp index b57e643a97..7242e0cff5 100644 --- a/xllm/core/distributed_runtime/llm_master.cpp +++ b/xllm/core/distributed_runtime/llm_master.cpp @@ -21,6 +21,7 @@ limitations under the License. #include #include +#include #include #include #include @@ -572,6 +573,10 @@ bool LLMMaster::unlink_p2p(const std::vector& remote_addrs) { return engine_->unlink_p2p(remote_addrs); } +bool LLMMaster::shutdown_remote_workers() { + return engine_->shutdown_remote_workers(); +} + LLMAssistantMaster::LLMAssistantMaster(const Options& options) : Master( options, @@ -589,6 +594,7 @@ LLMAssistantMaster::LLMAssistantMaster(const Options& options) } LLMAssistantMaster::~LLMAssistantMaster() { + stop(); // wait for the loop thread to finish if (loop_thread_.joinable()) { loop_thread_.join(); diff --git a/xllm/core/distributed_runtime/llm_master.h b/xllm/core/distributed_runtime/llm_master.h index ebcb31f94d..0f8fe58aff 100644 --- a/xllm/core/distributed_runtime/llm_master.h +++ b/xllm/core/distributed_runtime/llm_master.h @@ -100,6 +100,8 @@ class LLMMaster : public Master { void resume_scheduler(); bool is_scheduler_paused() const; + bool shutdown_remote_workers() override; + private: std::shared_ptr generate_request( std::string prompt, @@ -150,6 +152,7 @@ class LLMAssistantMaster : public Master { LLMAssistantMaster(const Options& options); ~LLMAssistantMaster(); void run() override; + void stop(); static void handle_signal(int signum) { running_ = false; } diff --git a/xllm/core/distributed_runtime/master.h b/xllm/core/distributed_runtime/master.h index 89a7831848..c8ca00c528 100644 --- a/xllm/core/distributed_runtime/master.h +++ b/xllm/core/distributed_runtime/master.h @@ -63,6 +63,8 @@ class Master { return false; } + virtual bool shutdown_remote_workers() { return false; } + MasterStatus get_master_status() const { return master_status_; } bool is_sleeping() const { return master_status_ != MasterStatus::WAKEUP; } diff --git a/xllm/core/distributed_runtime/remote_worker.cpp b/xllm/core/distributed_runtime/remote_worker.cpp index 26158665a9..2992411aa8 100644 --- a/xllm/core/distributed_runtime/remote_worker.cpp +++ b/xllm/core/distributed_runtime/remote_worker.cpp @@ -418,4 +418,13 @@ folly::SemiFuture RemoteWorker::stop_profile_async() { return future; } +folly::SemiFuture RemoteWorker::shutdown_async() { + folly::Promise promise; + auto future = promise.getSemiFuture(); + threadpool_.schedule([this, promise = std::move(promise)]() mutable { + promise.setValue(channel_->shutdown()); + }); + return future; +} + } // namespace xllm diff --git a/xllm/core/distributed_runtime/remote_worker.h b/xllm/core/distributed_runtime/remote_worker.h index d337989069..cc47dad0c1 100644 --- a/xllm/core/distributed_runtime/remote_worker.h +++ b/xllm/core/distributed_runtime/remote_worker.h @@ -157,6 +157,8 @@ class RemoteWorker : public WorkerClient { virtual folly::SemiFuture stop_profile_async() override; + virtual folly::SemiFuture shutdown_async() override; + private: DISALLOW_COPY_AND_ASSIGN(RemoteWorker); diff --git a/xllm/core/distributed_runtime/speculative_engine.cpp b/xllm/core/distributed_runtime/speculative_engine.cpp index 43d8505180..9963d1616b 100644 --- a/xllm/core/distributed_runtime/speculative_engine.cpp +++ b/xllm/core/distributed_runtime/speculative_engine.cpp @@ -293,4 +293,11 @@ bool SpeculativeEngine::unlink_cluster(const std::vector& cluster_ids, return engine_->unlink_cluster( cluster_ids, addrs, ports, src_dp_size, src_kv_split_size); }; + +bool SpeculativeEngine::shutdown_remote_workers() { + if (!dist_manager_) { + return true; + } + return dist_manager_->shutdown_remote_workers(); +} } // namespace xllm diff --git a/xllm/core/distributed_runtime/speculative_engine.h b/xllm/core/distributed_runtime/speculative_engine.h index 58b1c79883..7f15ed99d8 100644 --- a/xllm/core/distributed_runtime/speculative_engine.h +++ b/xllm/core/distributed_runtime/speculative_engine.h @@ -84,6 +84,8 @@ class SpeculativeEngine : public Engine { const int32_t src_dp_size, const int32_t src_kv_split_size = 1) override; + bool shutdown_remote_workers() override; + protected: SpeculativeEngine(const runtime::Options& options, bool use_draft_engine); diff --git a/xllm/core/distributed_runtime/vlm_engine.cpp b/xllm/core/distributed_runtime/vlm_engine.cpp index b5819efea0..829803be1c 100644 --- a/xllm/core/distributed_runtime/vlm_engine.cpp +++ b/xllm/core/distributed_runtime/vlm_engine.cpp @@ -445,6 +445,13 @@ void VLMEngine::setup_workers(const runtime::Options& options) { worker_clients_ = dist_manager_->get_worker_clients(); } +bool VLMEngine::shutdown_remote_workers() { + if (!dist_manager_) { + return true; + } + return dist_manager_->shutdown_remote_workers(); +} + std::vector VLMEngine::get_active_activation_memory() const { // call worker to get active activation memory std::vector> futures; diff --git a/xllm/core/distributed_runtime/vlm_engine.h b/xllm/core/distributed_runtime/vlm_engine.h index 22e66e83f7..75cff842f9 100644 --- a/xllm/core/distributed_runtime/vlm_engine.h +++ b/xllm/core/distributed_runtime/vlm_engine.h @@ -53,6 +53,8 @@ class VLMEngine : public Engine { // return the active activation memory std::vector get_active_activation_memory() const override; + bool shutdown_remote_workers() override; + private: bool init_model(); KVCacheCapacity estimate_kv_cache_capacity(); diff --git a/xllm/core/distributed_runtime/vlm_master.cpp b/xllm/core/distributed_runtime/vlm_master.cpp index 3c7fa3ca7c..9dbb7bd031 100644 --- a/xllm/core/distributed_runtime/vlm_master.cpp +++ b/xllm/core/distributed_runtime/vlm_master.cpp @@ -300,6 +300,10 @@ void VLMMaster::generate() { running_.store(false, std::memory_order_relaxed); } +bool VLMMaster::shutdown_remote_workers() { + return engine_->shutdown_remote_workers(); +} + std::shared_ptr VLMMaster::generate_request(std::string prompt, MMData mm_data, RequestParams sp, @@ -485,6 +489,7 @@ VLMAssistantMaster::VLMAssistantMaster(const Options& options) } VLMAssistantMaster::~VLMAssistantMaster() { + stop(); if (loop_thread_.joinable()) { loop_thread_.join(); } @@ -501,4 +506,8 @@ void VLMAssistantMaster::run() { }); } +void VLMAssistantMaster::stop() { + running_ = false; +} + } // namespace xllm diff --git a/xllm/core/distributed_runtime/vlm_master.h b/xllm/core/distributed_runtime/vlm_master.h index 22265959a9..6d5e2d9af0 100644 --- a/xllm/core/distributed_runtime/vlm_master.h +++ b/xllm/core/distributed_runtime/vlm_master.h @@ -82,6 +82,8 @@ class VLMMaster : public Master { // generate will run all requests, this is an blocking call void generate(); + bool shutdown_remote_workers() override; + int get_image_limit() { return options_.limit_image_per_prompt(); } private: @@ -129,6 +131,7 @@ class VLMAssistantMaster : public Master { explicit VLMAssistantMaster(const Options& options); ~VLMAssistantMaster(); void run() override; + void stop(); static void handle_signal(int signum) { running_ = false; } diff --git a/xllm/core/distributed_runtime/worker_service.cpp b/xllm/core/distributed_runtime/worker_service.cpp index ad0ad764d3..534b204f96 100644 --- a/xllm/core/distributed_runtime/worker_service.cpp +++ b/xllm/core/distributed_runtime/worker_service.cpp @@ -17,10 +17,14 @@ limitations under the License. #include #include #include +#include +#include #include #include #include +#include +#include #include #include "common/device_monitor.h" @@ -716,6 +720,20 @@ void WorkerService::StopProfile(::google::protobuf::RpcController* controller, return; } +void WorkerService::Shutdown(::google::protobuf::RpcController* controller, + const proto::Empty* req, + proto::Status* resp, + ::google::protobuf::Closure* done) { + (void)controller; + (void)req; + brpc::ClosureGuard done_guard(done); + resp->set_ok(true); + std::thread([]() { + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + kill(getpid(), SIGTERM); + }).detach(); +} + void WorkerService::ExecuteModel(::google::protobuf::RpcController* controller, const proto::ForwardInput* pb_forward_input, proto::ForwardOutput* pb_forward_output, diff --git a/xllm/core/distributed_runtime/worker_service.h b/xllm/core/distributed_runtime/worker_service.h index 01cbc62cdd..d042a836b4 100644 --- a/xllm/core/distributed_runtime/worker_service.h +++ b/xllm/core/distributed_runtime/worker_service.h @@ -150,6 +150,11 @@ class WorkerService : public proto::DistributeWorker { proto::Status* resp, ::google::protobuf::Closure* done) override; + void Shutdown(::google::protobuf::RpcController* controller, + const proto::Empty* req, + proto::Status* resp, + ::google::protobuf::Closure* done) override; + private: void step(ForwardInput& fwd_input, torch::Tensor& next_tokens, diff --git a/xllm/core/runtime/worker_client.cpp b/xllm/core/runtime/worker_client.cpp index 0759598e4e..8111e37f75 100644 --- a/xllm/core/runtime/worker_client.cpp +++ b/xllm/core/runtime/worker_client.cpp @@ -194,6 +194,10 @@ folly::SemiFuture WorkerClient::stop_profile_async() { return worker_->stop_profile_async(); } +folly::SemiFuture WorkerClient::shutdown_async() { + return folly::makeSemiFuture(true); +} + const torch::Device& WorkerClient::device() const { return worker_->device(); } folly::SemiFuture> diff --git a/xllm/core/runtime/worker_client.h b/xllm/core/runtime/worker_client.h index 8186527ccc..18a5c80de3 100644 --- a/xllm/core/runtime/worker_client.h +++ b/xllm/core/runtime/worker_client.h @@ -57,6 +57,8 @@ class WorkerClient { virtual folly::SemiFuture stop_profile_async(); + virtual folly::SemiFuture shutdown_async(); + virtual std::tuple estimate_kv_cache_capacity(); // allocate kv cache. blocking call diff --git a/xllm/proto/worker.proto b/xllm/proto/worker.proto index ed95104d06..bc2170d59e 100644 --- a/xllm/proto/worker.proto +++ b/xllm/proto/worker.proto @@ -371,4 +371,5 @@ service DistributeWorker { rpc UpdateWeights (UpdateWeightsRequest) returns (Status); rpc StartProfile (Empty) returns (Status); rpc StopProfile (Empty) returns (Status); + rpc Shutdown (Empty) returns (Status); } diff --git a/xllm/pybind/bind.cpp b/xllm/pybind/bind.cpp index 193027ccb8..526f35f7c3 100644 --- a/xllm/pybind/bind.cpp +++ b/xllm/pybind/bind.cpp @@ -144,6 +144,9 @@ PYBIND11_MODULE(xllm_export, m) { .def("options", &LLMMaster::options, py::call_guard()) + .def("shutdown_remote_workers", + &LLMMaster::shutdown_remote_workers, + py::call_guard()) .def( "build_sample_slots", [](const LLMMaster& self, @@ -195,6 +198,20 @@ PYBIND11_MODULE(xllm_export, m) { return "LLMMaster({})"_s.format(self.options()); }); + py::class_(m, "LLMAssistantMaster") + .def(py::init(), + py::arg("options"), + py::call_guard()) + .def("run", + &LLMAssistantMaster::run, + py::call_guard()) + .def("stop", + &LLMAssistantMaster::stop, + py::call_guard()) + .def("__repr__", [](const LLMAssistantMaster& self) { + return "LLMAssistantMaster({})"_s.format(self.options()); + }); + // 3. export SampleSlot py::class_(m, "SampleSlot") .def(py::init()) @@ -391,10 +408,27 @@ PYBIND11_MODULE(xllm_export, m) { .def("generate", &VLMMaster::generate, py::call_guard()) + .def("shutdown_remote_workers", + &VLMMaster::shutdown_remote_workers, + py::call_guard()) .def("__repr__", [](const VLMMaster& self) { return "VLMMaster({})"_s.format(self.options()); }); + py::class_(m, "VLMAssistantMaster") + .def(py::init(), + py::arg("options"), + py::call_guard()) + .def("run", + &VLMAssistantMaster::run, + py::call_guard()) + .def("stop", + &VLMAssistantMaster::stop, + py::call_guard()) + .def("__repr__", [](const VLMAssistantMaster& self) { + return "VLMAssistantMaster({})"_s.format(self.options()); + }); + // 12. export helpers m.def("get_model_backend", &ModelRegistry::get_model_backend, diff --git a/xllm/pybind/embedding.py b/xllm/pybind/embedding.py index db70ac7ecb..c7b2f4e13f 100644 --- a/xllm/pybind/embedding.py +++ b/xllm/pybind/embedding.py @@ -20,7 +20,7 @@ from . import utils from typing import Any, List, Optional, Union -from xllm_export import (LLMMaster, Options, RequestOutput, +from xllm_export import (LLMAssistantMaster, LLMMaster, Options, RequestOutput, RequestParams) from .errors import ValidationError @@ -44,6 +44,7 @@ def __init__( expert_parallel_degree: int = 0, enable_chunked_prefill: bool = True, enable_prefill_sp: bool = False, + master_node_addr: str = '', instance_role: str = 'DEFAULT', nnodes: int = 1, node_rank: int = 0, @@ -95,8 +96,12 @@ def __init__( options.expert_parallel_degree = expert_parallel_degree options.enable_chunked_prefill = enable_chunked_prefill options.enable_prefill_sp = enable_prefill_sp - free_port = utils.get_free_port() - options.master_node_addr = "127.0.0.1:" + str(free_port) + utils.require_multi_node_master_addr(nnodes, master_node_addr) + if master_node_addr: + options.master_node_addr = master_node_addr + else: + free_port = utils.get_free_port() + options.master_node_addr = "127.0.0.1:" + str(free_port) options.nnodes = nnodes options.node_rank = node_rank options.dp_size = dp_size @@ -114,10 +119,18 @@ def __init__( options.input_shm_size = input_shm_size options.output_shm_size = output_shm_size utils._configure_cpp_chat_template(use_cpp_chat_template, model_type) + self._is_driver = not utils.is_offline_worker_node(nnodes, node_rank) + if not self._is_driver: + self.master = LLMAssistantMaster(options) + self.master.run() + utils.block_offline_worker(self.master) + return self.master = LLMMaster(options) def finish(self) -> None: try: + if getattr(self, "_is_driver", True): + self.master.shutdown_remote_workers() #os.kill(os.getpid(), signal.SIGTERM) #os.kill(os.getpid(), signal.SIGKILL) utils.terminate_process(os.getpid()) @@ -130,6 +143,8 @@ def embedding( request_params: Optional[Union[RequestParams, List[RequestParams]]] = None, wait_for_schedule: bool = True, ) -> List[RequestOutput]: + if not getattr(self, "_is_driver", True): + raise RuntimeError("Only node_rank=0 can call embedding().") if request_params is None: request_params = RequestParams() if isinstance(inputs, str): diff --git a/xllm/pybind/llm.py b/xllm/pybind/llm.py index 62e1242f51..d71a8a956c 100644 --- a/xllm/pybind/llm.py +++ b/xllm/pybind/llm.py @@ -22,8 +22,8 @@ from . import utils from typing import Any, Callable, Dict, List, Optional, Sequence, Union -from xllm_export import (LLMMaster, VLMMaster, Options, RequestOutput, - RequestParams) +from xllm_export import (LLMAssistantMaster, LLMMaster, VLMAssistantMaster, + VLMMaster, Options, RequestOutput, RequestParams) from .errors import ValidationError from .params import ( BeamSearchParams, @@ -185,6 +185,7 @@ def __init__( options.expert_parallel_degree = expert_parallel_degree options.enable_chunked_prefill = enable_chunked_prefill options.enable_prefill_sp = enable_prefill_sp + utils.require_multi_node_master_addr(nnodes, master_node_addr) if master_node_addr: options.master_node_addr = master_node_addr else: @@ -216,6 +217,15 @@ def __init__( options.disable_log_stats = disable_log_stats options.enable_sleep_mode = enable_sleep_mode self._backend = backend + self._is_driver = not utils.is_offline_worker_node(nnodes, node_rank) + if not self._is_driver: + if backend == "vlm": + self.master = VLMAssistantMaster(options) + else: + self.master = LLMAssistantMaster(options) + self.master.run() + utils.block_offline_worker(self.master) + return if backend == "vlm": self.master = VLMMaster(options) else: @@ -223,6 +233,8 @@ def __init__( def finish(self) -> None: try: + if getattr(self, "_is_driver", True): + self.master.shutdown_remote_workers() #os.kill(os.getpid(), signal.SIGTERM) #os.kill(os.getpid(), signal.SIGKILL) utils.terminate_process(os.getpid()) @@ -279,6 +291,8 @@ def generate( use_tqdm: Union[bool, Callable[..., Any]] = True, **kwargs: Any, ) -> List[RequestOutput]: + if not getattr(self, "_is_driver", True): + raise RuntimeError("Only node_rank=0 can call generate().") request_params = kwargs.pop("request_params", None) if kwargs: unknown = ", ".join(kwargs.keys()) diff --git a/xllm/pybind/utils.py b/xllm/pybind/utils.py index 621bb0cd36..a8179db7ab 100644 --- a/xllm/pybind/utils.py +++ b/xllm/pybind/utils.py @@ -14,8 +14,11 @@ import json import os +import signal import socket -from typing import Dict, Optional, Tuple, Union +import sys +import time +from typing import Any, Dict, Optional, Tuple, Union import psutil import xllm_export @@ -51,6 +54,56 @@ def get_free_port() -> int: return port +def require_multi_node_master_addr( + nnodes: int, + master_node_addr: str, +) -> None: + if nnodes < 1: + raise ValueError("nnodes must be greater than or equal to 1.") + if nnodes <= 1: + return + if not master_node_addr: + raise ValueError( + "master_node_addr is required for multi-node offline inference." + ) + host, sep, port = master_node_addr.rpartition(":") + if not sep or not host or not port.isdigit(): + raise ValueError( + "master_node_addr must be in host:port format for " + "multi-node offline inference." + ) + if host in ("127.0.0.1", "localhost", "::1", "0.0.0.0", "::"): + raise ValueError( + "master_node_addr must be reachable by all nodes for " + "multi-node offline inference." + ) + + +def is_offline_worker_node(nnodes: int, node_rank: int) -> bool: + if node_rank < 0 or node_rank >= nnodes: + raise ValueError("node_rank must be in range [0, nnodes).") + return nnodes > 1 and node_rank > 0 + + +def block_offline_worker(master: Optional[Any] = None) -> None: + stop = False + + def _handle_signal(signum, frame) -> None: + nonlocal stop + if master is not None: + stop_master = getattr(master, "stop", None) + if callable(stop_master): + stop_master() + stop = True + + signal.signal(signal.SIGTERM, _handle_signal) + signal.signal(signal.SIGINT, _handle_signal) + + while not stop: + time.sleep(1) + sys.exit(0) + + def _read_json(path: str) -> Dict[str, object]: with open(path, "r", encoding="utf-8") as f: return json.load(f) diff --git a/xllm/pybind/vlm.py b/xllm/pybind/vlm.py index 8500addc8d..4cc58207bb 100644 --- a/xllm/pybind/vlm.py +++ b/xllm/pybind/vlm.py @@ -19,7 +19,7 @@ import threading from . import utils from typing import Any, Callable, Dict, List, Optional, Union -from xllm_export import (VLMMaster, Options, RequestOutput, +from xllm_export import (VLMAssistantMaster, VLMMaster, Options, RequestOutput, RequestParams, MMData) from .errors import ValidationError from .params import ( @@ -70,6 +70,7 @@ def __init__( expert_parallel_degree: int = 0, enable_chunked_prefill: bool = True, enable_prefill_sp: bool = False, + master_node_addr: str = '', instance_role: str = 'DEFAULT', transfer_listen_port: int = 26000, nnodes: int = 1, @@ -128,8 +129,12 @@ def __init__( options.expert_parallel_degree = expert_parallel_degree options.enable_chunked_prefill = enable_chunked_prefill options.enable_prefill_sp = enable_prefill_sp - free_port = utils.get_free_port() - options.master_node_addr = "127.0.0.1:" + str(free_port) + utils.require_multi_node_master_addr(nnodes, master_node_addr) + if master_node_addr: + options.master_node_addr = master_node_addr + else: + free_port = utils.get_free_port() + options.master_node_addr = "127.0.0.1:" + str(free_port) options.transfer_listen_port = transfer_listen_port options.nnodes = nnodes options.node_rank = node_rank @@ -151,10 +156,18 @@ def __init__( options.output_shm_size = output_shm_size options.disable_log_stats = disable_log_stats utils._configure_cpp_chat_template(use_cpp_chat_template, model_type) + self._is_driver = not utils.is_offline_worker_node(nnodes, node_rank) + if not self._is_driver: + self.master = VLMAssistantMaster(options) + self.master.run() + utils.block_offline_worker(self.master) + return self.master = VLMMaster(options) def finish(self) -> None: try: + if getattr(self, "_is_driver", True): + self.master.shutdown_remote_workers() #os.kill(os.getpid(), signal.SIGTERM) #os.kill(os.getpid(), signal.SIGKILL) utils.terminate_process(os.getpid()) @@ -177,6 +190,8 @@ def generate( use_tqdm: Union[bool, Callable[..., Any]] = True, **kwargs: Any, ) -> List[RequestOutput]: + if not getattr(self, "_is_driver", True): + raise RuntimeError("Only node_rank=0 can call generate().") from . import mm_utils prompts, mm_datas, image_urls = mm_utils.normalize_vllm_style_inputs(prompts) From 8a7d8805dc47b8910be584ef04e64a38571f8f79 Mon Sep 17 00:00:00 2001 From: huangweizhe1 Date: Thu, 18 Jun 2026 19:55:07 +0800 Subject: [PATCH 2/3] bugfix --- .../core/distributed_runtime/dist_manager.cpp | 2 +- xllm/core/distributed_runtime/dit_master.cpp | 12 ++++++---- xllm/core/distributed_runtime/dit_master.h | 1 + xllm/core/distributed_runtime/llm_master.cpp | 13 ++++++---- xllm/core/distributed_runtime/llm_master.h | 2 +- xllm/core/distributed_runtime/master.h | 2 ++ xllm/core/distributed_runtime/vlm_master.cpp | 12 +++++----- xllm/core/distributed_runtime/vlm_master.h | 2 +- xllm/pybind/bind.cpp | 8 +++---- xllm/pybind/embedding.py | 4 ++-- xllm/pybind/llm.py | 4 ++-- xllm/pybind/utils.py | 24 +------------------ xllm/pybind/vlm.py | 4 ++-- xllm/xllm.cpp | 3 +++ 14 files changed, 42 insertions(+), 51 deletions(-) diff --git a/xllm/core/distributed_runtime/dist_manager.cpp b/xllm/core/distributed_runtime/dist_manager.cpp index 96e1847d1d..aaf1944076 100644 --- a/xllm/core/distributed_runtime/dist_manager.cpp +++ b/xllm/core/distributed_runtime/dist_manager.cpp @@ -55,7 +55,7 @@ DistManager::~DistManager() { HealthCheckManager::instance().stop_health_check_thread(); XllmServer* collective_server = - ServerRegistry::get_instance().get_server(server_name_); + ServerRegistry::get_instance().try_get_server(server_name_); if (collective_server != nullptr) { collective_server->stop(); diff --git a/xllm/core/distributed_runtime/dit_master.cpp b/xllm/core/distributed_runtime/dit_master.cpp index 0f4be41fed..963f019239 100644 --- a/xllm/core/distributed_runtime/dit_master.cpp +++ b/xllm/core/distributed_runtime/dit_master.cpp @@ -165,10 +165,8 @@ DiTAssistantMaster::DiTAssistantMaster(const Options& options) } DiTAssistantMaster::~DiTAssistantMaster() { - // wait for the loop thread to finish - if (loop_thread_.joinable()) { - loop_thread_.join(); - } + running_ = false; + wait(); } void DiTAssistantMaster::run() { @@ -182,4 +180,10 @@ void DiTAssistantMaster::run() { }); } +void DiTAssistantMaster::wait() { + if (loop_thread_.joinable()) { + loop_thread_.join(); + } +} + } // namespace xllm diff --git a/xllm/core/distributed_runtime/dit_master.h b/xllm/core/distributed_runtime/dit_master.h index d1d51a885f..a04c3cc526 100644 --- a/xllm/core/distributed_runtime/dit_master.h +++ b/xllm/core/distributed_runtime/dit_master.h @@ -73,6 +73,7 @@ class DiTAssistantMaster : public Master { DiTAssistantMaster(const Options& options); ~DiTAssistantMaster(); void run() override; + void wait() override; static void handle_signal(int signum) { running_ = false; } diff --git a/xllm/core/distributed_runtime/llm_master.cpp b/xllm/core/distributed_runtime/llm_master.cpp index 7242e0cff5..7934b4b585 100644 --- a/xllm/core/distributed_runtime/llm_master.cpp +++ b/xllm/core/distributed_runtime/llm_master.cpp @@ -594,11 +594,8 @@ LLMAssistantMaster::LLMAssistantMaster(const Options& options) } LLMAssistantMaster::~LLMAssistantMaster() { - stop(); - // wait for the loop thread to finish - if (loop_thread_.joinable()) { - loop_thread_.join(); - } + running_ = false; + wait(); } void LLMAssistantMaster::run() { @@ -612,6 +609,12 @@ void LLMAssistantMaster::run() { }); } +void LLMAssistantMaster::wait() { + if (loop_thread_.joinable()) { + loop_thread_.join(); + } +} + // ============== Async RL training support: Pause/Resume ============== void LLMMaster::pause_scheduler(const std::string& mode) { LOG(INFO) << "LLMMaster: pausing scheduler (mode=" << mode << ")"; diff --git a/xllm/core/distributed_runtime/llm_master.h b/xllm/core/distributed_runtime/llm_master.h index 0f8fe58aff..ebdb463b4c 100644 --- a/xllm/core/distributed_runtime/llm_master.h +++ b/xllm/core/distributed_runtime/llm_master.h @@ -152,7 +152,7 @@ class LLMAssistantMaster : public Master { LLMAssistantMaster(const Options& options); ~LLMAssistantMaster(); void run() override; - void stop(); + void wait() override; static void handle_signal(int signum) { running_ = false; } diff --git a/xllm/core/distributed_runtime/master.h b/xllm/core/distributed_runtime/master.h index c8ca00c528..c81f7eecb7 100644 --- a/xllm/core/distributed_runtime/master.h +++ b/xllm/core/distributed_runtime/master.h @@ -65,6 +65,8 @@ class Master { virtual bool shutdown_remote_workers() { return false; } + virtual void wait() {} + MasterStatus get_master_status() const { return master_status_; } bool is_sleeping() const { return master_status_ != MasterStatus::WAKEUP; } diff --git a/xllm/core/distributed_runtime/vlm_master.cpp b/xllm/core/distributed_runtime/vlm_master.cpp index 9dbb7bd031..33a6a50598 100644 --- a/xllm/core/distributed_runtime/vlm_master.cpp +++ b/xllm/core/distributed_runtime/vlm_master.cpp @@ -489,10 +489,8 @@ VLMAssistantMaster::VLMAssistantMaster(const Options& options) } VLMAssistantMaster::~VLMAssistantMaster() { - stop(); - if (loop_thread_.joinable()) { - loop_thread_.join(); - } + running_ = false; + wait(); } void VLMAssistantMaster::run() { @@ -506,8 +504,10 @@ void VLMAssistantMaster::run() { }); } -void VLMAssistantMaster::stop() { - running_ = false; +void VLMAssistantMaster::wait() { + if (loop_thread_.joinable()) { + loop_thread_.join(); + } } } // namespace xllm diff --git a/xllm/core/distributed_runtime/vlm_master.h b/xllm/core/distributed_runtime/vlm_master.h index 6d5e2d9af0..86b0c237be 100644 --- a/xllm/core/distributed_runtime/vlm_master.h +++ b/xllm/core/distributed_runtime/vlm_master.h @@ -131,7 +131,7 @@ class VLMAssistantMaster : public Master { explicit VLMAssistantMaster(const Options& options); ~VLMAssistantMaster(); void run() override; - void stop(); + void wait() override; static void handle_signal(int signum) { running_ = false; } diff --git a/xllm/pybind/bind.cpp b/xllm/pybind/bind.cpp index 526f35f7c3..0e99bf4b9b 100644 --- a/xllm/pybind/bind.cpp +++ b/xllm/pybind/bind.cpp @@ -205,8 +205,8 @@ PYBIND11_MODULE(xllm_export, m) { .def("run", &LLMAssistantMaster::run, py::call_guard()) - .def("stop", - &LLMAssistantMaster::stop, + .def("wait", + &LLMAssistantMaster::wait, py::call_guard()) .def("__repr__", [](const LLMAssistantMaster& self) { return "LLMAssistantMaster({})"_s.format(self.options()); @@ -422,8 +422,8 @@ PYBIND11_MODULE(xllm_export, m) { .def("run", &VLMAssistantMaster::run, py::call_guard()) - .def("stop", - &VLMAssistantMaster::stop, + .def("wait", + &VLMAssistantMaster::wait, py::call_guard()) .def("__repr__", [](const VLMAssistantMaster& self) { return "VLMAssistantMaster({})"_s.format(self.options()); diff --git a/xllm/pybind/embedding.py b/xllm/pybind/embedding.py index c7b2f4e13f..8042823497 100644 --- a/xllm/pybind/embedding.py +++ b/xllm/pybind/embedding.py @@ -123,8 +123,8 @@ def __init__( if not self._is_driver: self.master = LLMAssistantMaster(options) self.master.run() - utils.block_offline_worker(self.master) - return + self.master.wait() + sys.exit(0) self.master = LLMMaster(options) def finish(self) -> None: diff --git a/xllm/pybind/llm.py b/xllm/pybind/llm.py index d71a8a956c..f85b04150c 100644 --- a/xllm/pybind/llm.py +++ b/xllm/pybind/llm.py @@ -224,8 +224,8 @@ def __init__( else: self.master = LLMAssistantMaster(options) self.master.run() - utils.block_offline_worker(self.master) - return + self.master.wait() + sys.exit(0) if backend == "vlm": self.master = VLMMaster(options) else: diff --git a/xllm/pybind/utils.py b/xllm/pybind/utils.py index a8179db7ab..813439fee1 100644 --- a/xllm/pybind/utils.py +++ b/xllm/pybind/utils.py @@ -14,11 +14,8 @@ import json import os -import signal import socket -import sys -import time -from typing import Any, Dict, Optional, Tuple, Union +from typing import Dict, Optional, Tuple, Union import psutil import xllm_export @@ -85,25 +82,6 @@ def is_offline_worker_node(nnodes: int, node_rank: int) -> bool: return nnodes > 1 and node_rank > 0 -def block_offline_worker(master: Optional[Any] = None) -> None: - stop = False - - def _handle_signal(signum, frame) -> None: - nonlocal stop - if master is not None: - stop_master = getattr(master, "stop", None) - if callable(stop_master): - stop_master() - stop = True - - signal.signal(signal.SIGTERM, _handle_signal) - signal.signal(signal.SIGINT, _handle_signal) - - while not stop: - time.sleep(1) - sys.exit(0) - - def _read_json(path: str) -> Dict[str, object]: with open(path, "r", encoding="utf-8") as f: return json.load(f) diff --git a/xllm/pybind/vlm.py b/xllm/pybind/vlm.py index 4cc58207bb..e5a2469413 100644 --- a/xllm/pybind/vlm.py +++ b/xllm/pybind/vlm.py @@ -160,8 +160,8 @@ def __init__( if not self._is_driver: self.master = VLMAssistantMaster(options) self.master.run() - utils.block_offline_worker(self.master) - return + self.master.wait() + sys.exit(0) self.master = VLMMaster(options) def finish(self) -> None: diff --git a/xllm/xllm.cpp b/xllm/xllm.cpp index 2d5b2ace58..545b9b813b 100644 --- a/xllm/xllm.cpp +++ b/xllm/xllm.cpp @@ -462,6 +462,9 @@ int run() { } else { master = std::make_unique(options); } + master->run(); + master->wait(); + return 0; } else { // master node master = create_master(model_config.backend(), options); From 4cde33b0be9f6812ad37fe707fc6ed20df618372 Mon Sep 17 00:00:00 2001 From: huangweizhe1 Date: Wed, 1 Jul 2026 14:54:02 +0800 Subject: [PATCH 3/3] docs: add comment and update docs. --- docs/en/getting_started/offline_service.md | 56 +++++++++++++++++++ docs/zh/getting_started/offline_service.md | 56 +++++++++++++++++++ xllm/core/distributed_runtime/comm_channel.h | 2 + .../core/distributed_runtime/dist_manager.cpp | 3 +- xllm/core/distributed_runtime/dist_manager.h | 3 + xllm/core/distributed_runtime/engine.h | 1 + xllm/core/distributed_runtime/master.h | 4 ++ xllm/core/distributed_runtime/remote_worker.h | 2 + .../distributed_runtime/worker_service.cpp | 2 +- .../core/distributed_runtime/worker_service.h | 2 + xllm/core/runtime/worker_client.h | 3 + xllm/proto/worker.proto | 3 + 12 files changed, 134 insertions(+), 3 deletions(-) diff --git a/docs/en/getting_started/offline_service.md b/docs/en/getting_started/offline_service.md index d5b3cacd2d..f11589611d 100644 --- a/docs/en/getting_started/offline_service.md +++ b/docs/en/getting_started/offline_service.md @@ -31,6 +31,62 @@ llm.finish() For LLM Beam Search, use `beam_width` as the switch. `top_logprobs` controls the top-k candidate count used for beam expansion at each decode step. If `top_logprobs` is left at its default value, xLLM uses `beam_width` as the top logprob count. Set `top_logprobs` to a value greater than `beam_width` when you want each beam to consider more candidate tokens. This beam-search top-k is different from the sampling cutoff parameter `top_k`. `best_of` is not the Beam Search switch, and this offline LLM guide does not use `num_return_sequences` to control the returned beams. +## Multi-Machine Offline Inference + +When a single machine does not have enough devices to load the model, you can use multi-machine offline inference to distribute the model across multiple machines. All machines run the same script and are differentiated by the `nnodes`, `node_rank`, and `master_node_addr` parameters: + +- The machine with `node_rank=0` is the **driver**, responsible for submitting inference requests and collecting results +- Machines with `node_rank>0` are **assistants**, which enter a wait loop after initialization and are scheduled by the driver + +### Parameters + +| Parameter | Description | +|-----------|-------------| +| `nnodes` | Total number of machines, defaults to 1 (single machine) | +| `node_rank` | Current machine index, range `[0, nnodes)`, driver is 0 | +| `master_node_addr` | Communication address of the driver machine in `host:port` format; must be reachable from all machines | +| `rank_tablefile` | Path to the HCCL collective communication topology file (required for multi-machine) | + +### Example + +Using 2 machines with 16 devices total, assuming the driver IP is `11.87.191.99`: + +**Machine 0 (driver):** + +```bash +python examples/generate.py \ + --model='/path/models/GLM-5-final-w8a8' \ + --devices='npu:0,npu:1,npu:2,npu:3,npu:4,npu:5,npu:6,npu:7' \ + --max_seqs_per_batch=300 \ + --max_memory_utilization=0.9 \ + --nnodes=2 \ + --node_rank=0 \ + --rank_tablefile=hccl_tools/hccl_2s_16p.json \ + --master_node_addr="11.87.191.99:19888" +``` + +**Machine 1 (assistant):** + +```bash +python examples/generate.py \ + --model='/path/models/GLM-5-final-w8a8' \ + --devices='npu:0,npu:1,npu:2,npu:3,npu:4,npu:5,npu:6,npu:7' \ + --max_seqs_per_batch=300 \ + --max_memory_utilization=0.9 \ + --nnodes=2 \ + --node_rank=1 \ + --rank_tablefile=hccl_tools/hccl_2s_16p.json \ + --master_node_addr="11.87.191.99:19888" +``` + +### Notes + +1. All machines must use the same model path, `devices` configuration, and `nnodes` value +2. `master_node_addr` must be the actual reachable IP of the driver; loopback addresses (`127.0.0.1`, `localhost`, etc.) are not allowed +3. `rank_tablefile` must be pre-generated via `hccl_tools` to describe the multi-machine collective communication topology +4. Assistant machines only need to construct the `LLM` object — there is no need to call `generate()` or define prompts; during `LLM` initialization, when `node_rank>0` is detected, the process automatically enters a wait loop and is shut down remotely when the driver calls `llm.finish()` +5. Machines can be started in any order; the framework waits internally for all machines to be ready + ## Embedding Generate embedding example: [:simple-github: https://github.com/jd-opensource/xllm/blob/main/examples/generate_embedding.py](https://github.com/jd-opensource/xllm/blob/main/examples/generate_embedding.py) diff --git a/docs/zh/getting_started/offline_service.md b/docs/zh/getting_started/offline_service.md index 9e68d57bca..54655e2a26 100644 --- a/docs/zh/getting_started/offline_service.md +++ b/docs/zh/getting_started/offline_service.md @@ -31,6 +31,62 @@ llm.finish() LLM Beam Search 使用 `beam_width` 作为开启参数。`top_logprobs` 控制每个解码步用于 beam 扩展的 top-k 候选数量。如果 `top_logprobs` 保持默认值，xLLM 会使用 `beam_width` 作为 top logprob 数量。如果希望每个 beam 考虑更多候选 token，可以将 `top_logprobs` 设置为大于 `beam_width` 的值。这里的 beam-search top-k 不同于采样截断参数 `top_k`。`best_of` 不是 Beam Search 开关，本文档也不使用 `num_return_sequences` 来控制 LLM 返回的 beam 数。 +## 多机离线推理 + +当单台机器的设备数量不足以加载模型时，可以使用多机离线推理将模型分布在多个节点上。所有节点运行相同的脚本，通过 `nnodes`、`node_rank` 和 `master_node_addr` 参数区分角色： + +- `node_rank=0` 的节点为驱动节点（driver），负责发起推理请求和收集结果 +- `node_rank>0` 的节点为辅助节点（assistant），启动后自动进入等待状态，由驱动节点调度执行 + +### 参数说明 + +| 参数 | 说明 | +|------|------| +| `nnodes` | 总节点数，默认为 1（单机） | +| `node_rank` | 当前节点编号，范围 `[0, nnodes)`，驱动节点为 0 | +| `master_node_addr` | 驱动节点的通信地址，格式为 `host:port`，所有节点必须能访问该地址 | +| `rank_tablefile` | HCCL 集合通信拓扑文件路径（多机场景必须提供） | + +### 使用示例 + +以 2 机 16 卡为例，假设驱动节点 IP 为 `11.87.191.99`： + +**节点 0（驱动节点）：** + +```bash +python examples/generate.py \ + --model='/path/models/GLM-5-final-w8a8' \ + --devices='npu:0,npu:1,npu:2,npu:3,npu:4,npu:5,npu:6,npu:7' \ + --max_seqs_per_batch=300 \ + --max_memory_utilization=0.9 \ + --nnodes=2 \ + --node_rank=0 \ + --rank_tablefile=hccl_tools/hccl_2s_16p.json \ + --master_node_addr="11.87.191.99:19888" +``` + +**节点 1（辅助节点）：** + +```bash +python examples/generate.py \ + --model='/path/models/GLM-5-final-w8a8' \ + --devices='npu:0,npu:1,npu:2,npu:3,npu:4,npu:5,npu:6,npu:7' \ + --max_seqs_per_batch=300 \ + --max_memory_utilization=0.9 \ + --nnodes=2 \ + --node_rank=1 \ + --rank_tablefile=hccl_tools/hccl_2s_16p.json \ + --master_node_addr="11.87.191.99:19888" +``` + +### 注意事项 + +1. 所有节点必须使用相同的模型路径、`devices` 配置和 `nnodes` 参数 +2. `master_node_addr` 必须为驱动节点的真实可达 IP，不能使用 `127.0.0.1`、`localhost` 等回环地址 +3. `rank_tablefile` 需预先通过 `hccl_tools` 生成，描述多机集合通信拓扑 +4. 辅助机器只需构造 `LLM` 对象即可，无需调用 `generate()` 或定义 prompts；`LLM` 初始化时检测到 `node_rank>0` 会自动进入等待循环，直到驱动节点调用 `llm.finish()` 时被远程关闭 +5. 各节点启动顺序无要求，框架内部会自动等待所有节点就绪 + ## Embedding 生成Embedding示例：[:simple-github: https://github.com/jd-opensource/xllm/blob/main/examples/generate_embedding.py](https://github.com/jd-opensource/xllm/blob/main/examples/generate_embedding.py) diff --git a/xllm/core/distributed_runtime/comm_channel.h b/xllm/core/distributed_runtime/comm_channel.h index 939ed8c1da..7f43536bcb 100644 --- a/xllm/core/distributed_runtime/comm_channel.h +++ b/xllm/core/distributed_runtime/comm_channel.h @@ -122,6 +122,8 @@ class CommChannel { virtual bool stop_profile(); + // Request the remote worker to terminate. Currently only used by + // multi-machine offline inference for coordinated shutdown of all machines. virtual bool shutdown(); protected: diff --git a/xllm/core/distributed_runtime/dist_manager.cpp b/xllm/core/distributed_runtime/dist_manager.cpp index aaf1944076..bd2914c94b 100644 --- a/xllm/core/distributed_runtime/dist_manager.cpp +++ b/xllm/core/distributed_runtime/dist_manager.cpp @@ -15,9 +15,8 @@ limitations under the License. #include "distributed_runtime/dist_manager.h" -#include - #include +#include #include "comm_channel.h" #include "common/health_check_manager.h" diff --git a/xllm/core/distributed_runtime/dist_manager.h b/xllm/core/distributed_runtime/dist_manager.h index a8dd057bf3..440e433220 100755 --- a/xllm/core/distributed_runtime/dist_manager.h +++ b/xllm/core/distributed_runtime/dist_manager.h @@ -32,6 +32,9 @@ class DistManager { return worker_clients_; }; + // Shut down all remote (non-local) worker processes via the Shutdown RPC. + // Currently only used by multi-machine offline inference. Sends shutdown + // only to workers on other machines; skips local ranks. Idempotent. bool shutdown_remote_workers(); private: diff --git a/xllm/core/distributed_runtime/engine.h b/xllm/core/distributed_runtime/engine.h index efbfa4f716..dec8786940 100644 --- a/xllm/core/distributed_runtime/engine.h +++ b/xllm/core/distributed_runtime/engine.h @@ -171,6 +171,7 @@ class Engine { return false; }; + // Shut down remote workers (multi-machine offline inference only). virtual bool shutdown_remote_workers() { return true; }; // XTensor mode: get GlobalXTensor offsets for allocated blocks diff --git a/xllm/core/distributed_runtime/master.h b/xllm/core/distributed_runtime/master.h index c81f7eecb7..ee15ed0ff4 100644 --- a/xllm/core/distributed_runtime/master.h +++ b/xllm/core/distributed_runtime/master.h @@ -63,6 +63,10 @@ class Master { return false; } + // Shut down all remote worker machines. Currently only used by + // multi-machine offline inference: the driver (node_rank=0) calls this + // inside LLM.finish() to terminate assistant machines so all processes + // exit cleanly. virtual bool shutdown_remote_workers() { return false; } virtual void wait() {} diff --git a/xllm/core/distributed_runtime/remote_worker.h b/xllm/core/distributed_runtime/remote_worker.h index cc47dad0c1..7b48ba53b5 100644 --- a/xllm/core/distributed_runtime/remote_worker.h +++ b/xllm/core/distributed_runtime/remote_worker.h @@ -157,6 +157,8 @@ class RemoteWorker : public WorkerClient { virtual folly::SemiFuture stop_profile_async() override; + // Asynchronously request this remote worker to shut down + // (multi-machine offline inference only). virtual folly::SemiFuture shutdown_async() override; private: diff --git a/xllm/core/distributed_runtime/worker_service.cpp b/xllm/core/distributed_runtime/worker_service.cpp index 534b204f96..20f3bcd363 100644 --- a/xllm/core/distributed_runtime/worker_service.cpp +++ b/xllm/core/distributed_runtime/worker_service.cpp @@ -18,8 +18,8 @@ limitations under the License. #include #include #include -#include #include +#include #include #include diff --git a/xllm/core/distributed_runtime/worker_service.h b/xllm/core/distributed_runtime/worker_service.h index d042a836b4..f7aed13e00 100644 --- a/xllm/core/distributed_runtime/worker_service.h +++ b/xllm/core/distributed_runtime/worker_service.h @@ -150,6 +150,8 @@ class WorkerService : public proto::DistributeWorker { proto::Status* resp, ::google::protobuf::Closure* done) override; + // Shutdown RPC handler for multi-machine offline inference. Sends SIGTERM to + // self after a short delay (to allow the RPC response to be sent first). void Shutdown(::google::protobuf::RpcController* controller, const proto::Empty* req, proto::Status* resp, diff --git a/xllm/core/runtime/worker_client.h b/xllm/core/runtime/worker_client.h index 18a5c80de3..2c057b8fe9 100644 --- a/xllm/core/runtime/worker_client.h +++ b/xllm/core/runtime/worker_client.h @@ -57,6 +57,9 @@ class WorkerClient { virtual folly::SemiFuture stop_profile_async(); + // Request the worker to terminate. Currently only used by multi-machine + // offline inference. Base implementation is a no-op; RemoteWorker overrides + // to send the Shutdown RPC. virtual folly::SemiFuture shutdown_async(); virtual std::tuple estimate_kv_cache_capacity(); diff --git a/xllm/proto/worker.proto b/xllm/proto/worker.proto index bc2170d59e..09e071b572 100644 --- a/xllm/proto/worker.proto +++ b/xllm/proto/worker.proto @@ -371,5 +371,8 @@ service DistributeWorker { rpc UpdateWeights (UpdateWeightsRequest) returns (Status); rpc StartProfile (Empty) returns (Status); rpc StopProfile (Empty) returns (Status); + // Gracefully terminate a remote worker process. Currently only used by + // multi-machine offline inference: the driver node (node_rank=0) calls this + // on every remote worker after generation completes so all processes exit. rpc Shutdown (Empty) returns (Status); }