refactor turbomind (2/N) (#2818)

InternLM · Nov 29, 2024 · 4ede631 · 4ede631
1 parent 0b6dd1f
commit 4ede631
Show file tree

Hide file tree

Showing 13 changed files with 370 additions and 696 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -304,7 +304,7 @@ link_directories(
 
 # add_subdirectory(3rdparty)
 add_subdirectory(src)
-add_subdirectory(examples)
+# add_subdirectory(examples)
 
 if(BUILD_TEST)
     add_subdirectory(tests/csrc)

diff --git a/lmdeploy/turbomind/turbomind.py b/lmdeploy/turbomind/turbomind.py
@@ -358,12 +358,10 @@ def _forward_callback(self, result, ctx):
         self.que.put((False, result))
 
     def _forward_thread(self, inputs):
-        instance_comm = self.tm_model.model_comm.create_instance_comm(
-            self.gpu_count)
 
         def _func():
             try:
-                output = self.model_inst.forward(inputs, instance_comm)
+                output = self.model_inst.forward(inputs)
             except Exception as e:
                 logger.error(f'unhandled exception: {e}')
                 self.que.put((-1, None))
@@ -377,12 +375,10 @@ def _async_forward_callback(self, result, ctx, que: LifoQueue):
         que.put((False, result))
 
     def _async_forward_thread(self, inputs, que: LifoQueue):
-        instance_comm = self.tm_model.model_comm.create_instance_comm(
-            self.gpu_count)
 
         def _func():
             try:
-                output = self.model_inst.forward(inputs, instance_comm)
+                output = self.model_inst.forward(inputs)
             except Exception as e:
                 logger.error(f'unhandled exception: {e}')
                 que.put((-1, None))

diff --git a/src/turbomind/models/llama/LlamaBatch.h b/src/turbomind/models/llama/LlamaBatch.h
@@ -12,7 +12,6 @@
 #include "src/turbomind/utils/allocator.h"
 #include "src/turbomind/utils/cublasMMWrapper.h"
 #include "src/turbomind/utils/cuda_utils.h"
-#include "src/turbomind/utils/instance_comm.h"
 #include <condition_variable>
 #include <curand_kernel.h>
 #include <mutex>
@@ -32,8 +31,7 @@ struct SharedState {
 };
 
 struct Control {
-    AbstractInstanceComm* comm;
-    Request::Callback     callback;
+    Request::Callback callback;
 };
 
 struct BatchState {

diff --git a/src/turbomind/models/llama/LlamaV2.h b/src/turbomind/models/llama/LlamaV2.h
@@ -21,6 +21,9 @@
 
 #pragma once
 
+#include <limits>
+#include <unordered_map>
+
 #include "src/turbomind/layers/DynamicDecodeLayer.h"
 #include "src/turbomind/models/llama/Barrier.h"
 #include "src/turbomind/models/llama/LlamaBatch.h"
@@ -31,10 +34,7 @@
 #include "src/turbomind/models/llama/unified_decoder.h"
 #include "src/turbomind/utils/allocator.h"
 #include "src/turbomind/utils/cublasMMWrapper.h"
-#include "src/turbomind/utils/instance_comm.h"
 #include "src/turbomind/utils/nccl_utils.h"
-#include <limits>
-#include <unordered_map>
 
 namespace turbomind {