Fix the lora adapter when lora path is none (#4799)

Co-authored-by: Beichen Ma <[email protected]>

Author: Qiaolin-Yu

python/sglang/srt/lora/lora_manager.py

@@ -133,10 +133,6 @@ def prepare_lora_batch(self, forward_batch: ForwardBatch):
         assert len(cur_uids) <= self.max_loras_per_batch
         self.memory_pool.prepare_lora_batch(cur_uids, self.loras)
 
-        # FIXME: Handle lora uid with None more safely
-        if cur_uids == set([None]):
-            return
-
         # set up batch info shared by all lora moruldes
         bs = forward_batch.batch_size
         seg_lens = (

python/sglang/srt/lora/mem_pool.py

@@ -163,7 +163,7 @@ def load_lora_weight_to_buffer(
         if uid is None:
             for i in range(self.num_layer):
                 for k in self.A_buffer.keys():
-                    self.A_buffer[k][i][buffer_id] *= 0
+                    self.A_buffer[k][i][buffer_id] = 0
             return
 
         assert lora_adapter is not None

test/srt/models/lora/test_lora.py

@@ -96,6 +96,11 @@ def inference(self, prompts, lora_set, tp_size, torch_dtype, max_new_tokens):
             srt_outputs = srt_runner.forward(
                 prompts, max_new_tokens=max_new_tokens, lora_paths=batch_lora_paths
             )
+            srt_outputs_lora_path_none = srt_runner.forward(
+                prompts,
+                max_new_tokens=max_new_tokens,
+                lora_paths=[None] * len(prompts),
+            )
 
         with HFRunner(
             base_path, torch_dtype=torch_dtype, model_type="generation"
@@ -169,18 +174,20 @@ def inference(self, prompts, lora_set, tp_size, torch_dtype, max_new_tokens):
         print(f"{srt_outputs.output_strs=}")
         print(f"{hf_no_lora_outputs.output_strs=}")
         print(f"{srt_no_lora_outputs.output_strs=}")
+        print(f"{srt_outputs_lora_path_none.output_strs=}")
         for i in range(len(prompts)):
             assert srt_outputs.output_strs[i].strip(" ") == hf_outputs.output_strs[i], (
                 srt_outputs.output_strs[i].strip(" "),
                 hf_outputs.output_strs[i],
             )
-            # assert (
-            #     srt_no_lora_outputs.output_strs[i].strip(" ")
-            #     == hf_no_lora_outputs.output_strs[i]
-            # ), (
-            #     srt_no_lora_outputs.output_strs[i].strip(" "),
-            #     hf_no_lora_outputs.output_strs[i],
-            # )
+            assert (
+                srt_no_lora_outputs.output_strs[i].strip(" ")
+                == hf_no_lora_outputs.output_strs[i]
+            ), (
+                srt_no_lora_outputs.output_strs[i].strip(" "),
+                hf_no_lora_outputs.output_strs[i],
+            )
+            assert srt_outputs_lora_path_none == srt_no_lora_outputs
 
     def serving(self, prompts, lora_set, tp_size, torch_dtype, max_new_tokens):
         print("=================== testing serving =======================")
@@ -257,7 +264,7 @@ def base_inference(self, prompts, lora_set, tp_size, torch_dtype, max_new_tokens
             srt_no_lora_logprobs = torch.Tensor(
                 srt_no_lora_outputs.top_input_logprobs[i]
             )
-            srt_logprobs = torch.uensor(srt_outputs.top_input_logprobs[i])
+            srt_logprobs = torch.Tensor(srt_outputs.top_input_logprobs[i])
             print("max_diff", torch.max(abs(srt_no_lora_logprobs - srt_logprobs)))
 
         print(f"{srt_no_lora_outputs.output_strs=}")
@@ -280,7 +287,7 @@ def test_all(self):
                 tp_size = 1
                 max_new_tokens = 32
                 self.inference(PROMPTS, lora_set, tp_size, torch_dtype, max_new_tokens)
-                # self.serving(PROMPTS, lora_set, tp_size, torch_dtype, max_new_tokens)
+                self.serving(PROMPTS, lora_set, tp_size, torch_dtype, max_new_tokens)
                 # self.base_inference(
                 #     PROMPTS, lora_set, tp_size, torch_dtype, max_new_tokens
                 # )