fix vllm_module InferenceMemoryManager args error (#61)

SeaOfOcean · web-flow · commit fd5d6cd3efe4 · 2024-09-06T13:37:16.000+08:00
diff --git a/chatlearn/models/vllm_module.py b/chatlearn/models/vllm_module.py
@@ -498,7 +498,7 @@ def model_setup(self):
             if InferenceMemoryManager is None:
                 raise Exception("Import InferenceMemoryManager failed, you may need to set right Megatron path first.")
             self._memory_manager = InferenceMemoryManager(
-                self.model, self.name, self.timers,
+                self.model,
                 self.runtime_args.bucket_size_mb_in_memory_manager,
             )
             self.offload()

Original file line number	Diff line number	Diff line change
`@@ -498,7 +498,7 @@ def model_setup(self):`
`498`	`498`	`if InferenceMemoryManager is None:`
`499`	`499`	`raise Exception("Import InferenceMemoryManager failed, you may need to set right Megatron path first.")`
`500`	`500`	`self._memory_manager = InferenceMemoryManager(`
`501`		`- self.model, self.name, self.timers,`
	`501`	`+ self.model,`
`502`	`502`	`self.runtime_args.bucket_size_mb_in_memory_manager,`
`503`	`503`	`)`
`504`	`504`	`self.offload()`