From cc46aa9828f3ff3223d895cd92d124cdccc06dc3 Mon Sep 17 00:00:00 2001
From: DeEMO <yzzxrx@gmail.com>
Date: Thu, 15 May 2025 13:57:16 +0800
Subject: [PATCH] update rank and n_world

Signed-off-by: DeEMO <yzzxrx@gmail.com>
---
 common/common.cpp | 24 ++++++++++++++++++++++++
 include/llama.h   |  6 +++++-
 src/llama.cpp     |  9 +++++++++
 3 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/common/common.cpp b/common/common.cpp
index 991b34a5..88b00075 100644
--- a/common/common.cpp
+++ b/common/common.cpp
@@ -1702,6 +1702,30 @@ struct llama_init_result llama_init_from_gpt_params(gpt_params & params) {
             }
         }
 
+        //update rank and n_world for consistency
+        uint32_t update_rank = 0;
+        uint32_t update_n_world = 1;
+        std::vector<uint32_t> n_layer_window_temp = {n_layer_window[0]};
+        std::vector<uint32_t> n_gpu_layers_temp = {n_gpu_layers[0]};
+        for(auto i=1; i<n_world; i++) {
+            if(n_layer_window[i] <= 0 ){
+                continue;
+            }
+            if(i <= my_rank){
+                update_rank++;
+            }
+            update_n_world++;
+            n_layer_window_temp.push_back(n_layer_window[i]);
+            n_gpu_layers_temp.push_back(n_gpu_layers[i]);
+        }
+        memset(n_layer_window, 0, n_world * sizeof(uint32_t));
+        memset(n_gpu_layers, 0, n_world * sizeof(uint32_t));
+        for (auto i=0; i<update_n_world; i++) {
+            n_layer_window[i] = n_layer_window_temp[i];
+            n_gpu_layers[i] = n_gpu_layers_temp[i];
+        }
+        llama_update_context_with_rankworld(lctx, update_rank, update_n_world);
+
         // update n_layer_window and n_gpu_layers
         std::copy(std::begin(n_layer_window), std::end(n_layer_window), params.n_layer_window);
         std::copy(std::begin(n_layer_window), std::end(n_layer_window), cparams.n_layer_window);
diff --git a/include/llama.h b/include/llama.h
index 515dfd93..083a3a8f 100644
--- a/include/llama.h
+++ b/include/llama.h
@@ -462,7 +462,11 @@ extern "C" {
               struct llama_model_loader * ml,
               struct llama_model        * model,
               struct llama_model_params   params);
-
+    
+    LLAMA_API void llama_update_context_with_rankworld(struct llama_context * ctx,
+                                                                     uint32_t rank,
+                                                                     uint32_t n_world);
+    
     LLAMA_API struct llama_context * llama_new_context_with_model(
                      struct llama_model * model,
             struct llama_context_params   params);
diff --git a/src/llama.cpp b/src/llama.cpp
index 8cc37213..cfee40b6 100644
--- a/src/llama.cpp
+++ b/src/llama.cpp
@@ -20475,6 +20475,15 @@ void llama_free_sockets(struct llama_context * ctx, char ** msg) {
     }
 }
 
+void llama_update_context_with_rankworld(struct llama_context * ctx,
+                                                       uint32_t rank,
+                                                       uint32_t n_world) {
+    if(ctx) {
+        ctx->cparams.rank = rank;
+        ctx->cparams.n_world = n_world;
+    }
+}
+
 struct llama_context * llama_new_context_with_model(
                  struct llama_model * model,
         struct llama_context_params   params) {