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/__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..7f43536bcb 100644 --- a/xllm/core/distributed_runtime/comm_channel.h +++ b/xllm/core/distributed_runtime/comm_channel.h @@ -122,6 +122,10 @@ 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: 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..bd2914c94b 100644 --- a/xllm/core/distributed_runtime/dist_manager.cpp +++ b/xllm/core/distributed_runtime/dist_manager.cpp @@ -15,6 +15,7 @@ limitations under the License. #include "distributed_runtime/dist_manager.h" +#include #include #include "comm_channel.h" @@ -47,11 +48,13 @@ DistManager::DistManager(const runtime::Options& options) } DistManager::~DistManager() { + shutdown_remote_workers(); + // Stop health check 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(); @@ -177,6 +180,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 +340,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..440e433220 100755 --- a/xllm/core/distributed_runtime/dist_manager.h +++ b/xllm/core/distributed_runtime/dist_manager.h @@ -32,6 +32,11 @@ 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: DISALLOW_COPY_AND_ASSIGN(DistManager); void setup_multi_node_workers(const runtime::Options& options, @@ -57,5 +62,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/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/engine.h b/xllm/core/distributed_runtime/engine.h index ee7419b168..dec8786940 100644 --- a/xllm/core/distributed_runtime/engine.h +++ b/xllm/core/distributed_runtime/engine.h @@ -171,6 +171,9 @@ 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 // 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..7934b4b585 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,10 +594,8 @@ LLMAssistantMaster::LLMAssistantMaster(const Options& options) } LLMAssistantMaster::~LLMAssistantMaster() { - // wait for the loop thread to finish - if (loop_thread_.joinable()) { - loop_thread_.join(); - } + running_ = false; + wait(); } void LLMAssistantMaster::run() { @@ -606,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 ebcb31f94d..ebdb463b4c 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 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 89a7831848..ee15ed0ff4 100644 --- a/xllm/core/distributed_runtime/master.h +++ b/xllm/core/distributed_runtime/master.h @@ -63,6 +63,14 @@ 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() {} + 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..7b48ba53b5 100644 --- a/xllm/core/distributed_runtime/remote_worker.h +++ b/xllm/core/distributed_runtime/remote_worker.h @@ -157,6 +157,10 @@ 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: 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..33a6a50598 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,9 +489,8 @@ VLMAssistantMaster::VLMAssistantMaster(const Options& options) } VLMAssistantMaster::~VLMAssistantMaster() { - if (loop_thread_.joinable()) { - loop_thread_.join(); - } + running_ = false; + wait(); } void VLMAssistantMaster::run() { @@ -501,4 +504,10 @@ void VLMAssistantMaster::run() { }); } +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 22265959a9..86b0c237be 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 wait() override; 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..20f3bcd363 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..f7aed13e00 100644 --- a/xllm/core/distributed_runtime/worker_service.h +++ b/xllm/core/distributed_runtime/worker_service.h @@ -150,6 +150,13 @@ 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, + ::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..2c057b8fe9 100644 --- a/xllm/core/runtime/worker_client.h +++ b/xllm/core/runtime/worker_client.h @@ -57,6 +57,11 @@ 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(); // allocate kv cache. blocking call diff --git a/xllm/proto/worker.proto b/xllm/proto/worker.proto index ed95104d06..09e071b572 100644 --- a/xllm/proto/worker.proto +++ b/xllm/proto/worker.proto @@ -371,4 +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); } diff --git a/xllm/pybind/bind.cpp b/xllm/pybind/bind.cpp index 193027ccb8..0e99bf4b9b 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("wait", + &LLMAssistantMaster::wait, + 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("wait", + &VLMAssistantMaster::wait, + 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..8042823497 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() + self.master.wait() + sys.exit(0) 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..f85b04150c 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() + self.master.wait() + sys.exit(0) 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..813439fee1 100644 --- a/xllm/pybind/utils.py +++ b/xllm/pybind/utils.py @@ -51,6 +51,37 @@ 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 _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..e5a2469413 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() + self.master.wait() + sys.exit(0) 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) 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);