Repository: sgl-project/sglang
Pull request #4274: [Feature] Support Tensor Parallelism and Weight Slicing for Lora
Author: login: aoshen524
Created: 2025-03-10T20:45:11Z | Merged: 2025-03-19T03:33:08Z

## Motivation

https://github.com/sgl-project/sglang/issues/3414 reports issues regarding limited model support compared to test_generation_models.py. This PR introduces tensor parallelism and weight slicing for LoRA, alongside additional improvements to testing and functionality.

## Modifications


- [x] Implemented tensor parallelism support in LoRA, allowing efficient distribution of computations across multiple devices.
 - [x] Introduced LoRA weight slicing and refactor memory pool to facilitate distributed inference, optimizing memory usage and performance.
 - [x] Remove useless code for CPU-GPU weight transmission.
 - [x] Add error handling when num available gpu < num needed
 
 checklist:

- [ ] Remove tensor.contiguous() used in GPU

### Throughput for LLaMA 7B with Triton Backend in 4 x 3090-24GB

*LoRA Config: wissdw/4r2a_llama_hf*

|                | TP Size = 1 | TP Size = 2 | TP Size = 4 |
|----------------|-------------|-------------|-------------|
| **Use LoRA**   | 800 tok/s   | 725 tok/s   | 717 tok/s  |
| **No LoRA**    | 793 tok/s   | 783 tok/s   | 765 tok/s  |



Local CI test result for test/srt/models/lora/test_lora_tp.py
![image](https://github.com/user-attachments/assets/ce29d9e9-1236-49e6-a8bc-bb83eb23555d)



Diff:
diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml
index 36983115d9fc..a652ec3cda1f 100644
--- a/.github/workflows/pr-test.yml
+++ b/.github/workflows/pr-test.yml
@@ -127,6 +127,12 @@ jobs:
           cd test/srt
           python3 test_mla_tp.py
 
+      - name: Test lora tensor parallelism (TP=2)
+        timeout-minutes: 10
+        run: |
+          cd test/srt/models/lora
+          python3 test_lora_tp.py
+
   performance-test-1-gpu-part-1:
     if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') &&
         github.event.pull_request.draft == false
diff --git a/benchmark/lora/launch_server.py b/benchmark/lora/launch_server.py
index 418155dbf5b5..33ad848a6c41 100644
--- a/benchmark/lora/launch_server.py
+++ b/benchmark/lora/launch_server.py
@@ -22,7 +22,10 @@ def launch_server(args):
     cmd += f"--disable-radix --disable-cuda-graph "
     cmd += f"--max-loras-per-batch {args.max_loras_per_batch} "
     cmd += f"--max-running-requests {args.max_running_requests} "
-    cmd += f"--lora-backend {args.lora_backend}"
+    cmd += f"--lora-backend {args.lora_backend} "
+    cmd += f"--tp-size {args.tp_size} "
+    if args.disable_custom_all_reduce:
+        cmd += "--disable-custom-all-reduce"
     print(cmd)
     os.system(cmd)
 
@@ -48,6 +51,18 @@ def launch_server(args):
         type=str,
         default="triton",
     )
+    parser.add_argument(
+        "--tp-size",
+        type=int,
+        default=1,
+        help="Tensor parallel size for distributed inference",
+    )
+    # disable_custom_all_reduce
+    parser.add_argument(
+        "--disable-custom-all-reduce",
+        action="store_true",
+        help="Disable custom all reduce when device does not support p2p communication",
+    )
     args = parser.parse_args()
 
     launch_server(args)
diff --git a/python/sglang/srt/layers/linear.py b/python/sglang/srt/layers/linear.py
index 85748fa7434e..a5bb6281dded 100644
--- a/python/sglang/srt/layers/linear.py
+++ b/python/sglang/srt/layers/linear.py
@@ -782,6 +782,8 @@ def __init__(
         else:
             self.num_kv_heads = divide(self.total_num_kv_heads, tp_size)
             self.num_kv_head_replicas = 1
+        self.q_proj_shard_size = self.num_heads * self.head_size
+        self.kv_proj_shard_size = self.num_kv_heads * self.head_size
         input_size = self.hidden_size
         output_size = (
             (self.num_heads + 2 * self.num_kv_heads) * tp_size * self.head_size
diff --git a/python/sglang/srt/lora/layers.py b/python/sglang/srt/lora/layers.py
index 558fffa90270..7d4f560a00db 100644
--- a/python/sglang/srt/lora/layers.py
+++ b/python/sglang/srt/lora/layers.py
@@ -1,3 +1,5 @@
+from typing import List, Tuple
+
 import torch
 from torch import nn
 
@@ -38,8 +40,22 @@ def forward(self, x: torch.Tensor):
     def set_lora_info(self, *args):
         pass
 
+    def slice_lora_a_weights(self, A: torch.Tensor, tp_rank: int):
+        pass
+
+    def slice_lora_b_weights(self, B: torch.Tensor, tp_rank: int):
+        pass
+
 
 class VocabParallelEmbeddingWithLoRA(BaseLayerWithLoRA):
+    """
+    Vocab parallel embedding layer with support for LoRA (Low-Rank Adaptation).
+
+    Note: The current version does not yet implement the LoRA functionality.
+    This class behaves exactly the same as the base VocabParallelEmbedding.
+    Future versions will integrate LoRA functionality to support efficient parameter fine-tuning.
+    """
+
     def __init__(
         self,
         base_layer: VocabParallelEmbedding,
@@ -101,6 +117,16 @@ def forward(self, input_: torch.Tensor):
         output_bias = self.base_layer.bias if self.base_layer.skip_bias_add else None
         return output, output_bias
 
+    def slice_lora_a_weights(self, A: torch.Tensor, tp_rank: int):
+        return A
+
+    def slice_lora_b_weights(self, B: torch.Tensor, tp_rank: int):
+        shard_size = self.base_layer.output_partition_sizes[0]
+        start_idx = tp_rank * shard_size
+        end_idx = (tp_rank + 1) * shard_size
+        B = B[start_idx:end_idx, :]
+        return B
+
 
 class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):
     def __init__(
@@ -120,6 +146,7 @@ def set_lora_info(
         self.set_lora = True
         self.A_buffer_gate_up = A_buffer
         if self.lora_backend.fuse_stacked_lora_b:
+            # TODO: avoid using contiguous() in GPU.
             # B_buffer_gate_up: (num_lora, 2 * output_dim, r)
             self.B_buffer_gate_up = torch.cat(
                 (B_buffer[0], B_buffer[1]), dim=-2
@@ -142,6 +169,16 @@ def apply_lora(self, base_output: torch.Tensor, x: torch.Tensor) -> torch.Tensor
             else base_output + lora_output * self.scaling
         )
 
+    def slice_lora_a_weights(self, A: torch.Tensor, tp_rank: int):
+        return A
+
+    def slice_lora_b_weights(self, B: torch.Tensor, tp_rank: int):
+        # Since the outputs for both gate and up are identical, we use a random one.
+        shard_size = self.base_layer.output_partition_sizes[0]
+        start_idx = tp_rank * shard_size
+        end_idx = (tp_rank + 1) * shard_size
+        return B[:, start_idx:end_idx, :]
+
 
 class QKVParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):
     def init__(
@@ -210,6 +247,27 @@ def apply_lora(self, base_output: torch.Tensor, x: torch.Tensor) -> torch.Tensor
             else base_output + lora_output * self.scaling
         )
 
+    def slice_lora_a_weights(self, A: torch.Tensor, tp_rank: int):
+        return A
+
+    def slice_lora_b_weights(
+        self, B: List[torch.Tensor], tp_rank: int
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        B_q, B_kv = B
+        base_layer = self.base_layer
+        q_proj_shard_size = base_layer.q_proj_shard_size
+        kv_proj_shard_size = base_layer.kv_proj_shard_size
+        num_kv_head_replicas = base_layer.num_kv_head_replicas
+
+        q_start_idx = q_proj_shard_size * tp_rank
+        q_end_idx = q_start_idx + q_proj_shard_size
+
+        kv_shard_id = tp_rank // num_kv_head_replicas
+        kv_start_idx = kv_proj_shard_size * kv_shard_id
+        kv_end_idx = kv_start_idx + kv_proj_shard_size
+
+        return B_q[q_start_idx:q_end_idx, :], B_kv[:, kv_start_idx:kv_end_idx, :]
+
 
 class RowParallelLinearWithLoRA(BaseLayerWithLoRA):
     def __init__(
@@ -274,6 +332,16 @@ def forward(self, input_: torch.Tensor):
             output_bias = self.base_layer.bias
         return output, output_bias
 
+    def slice_lora_a_weights(self, A: torch.Tensor, tp_rank: int):
+        shard_size = self.base_layer.input_size_per_partition
+        start_idx = tp_rank * shard_size
+        end_idx = (tp_rank + 1) * shard_size
+        A = A[:, start_idx:end_idx].contiguous()
+        return A
+
+    def slice_lora_b_weights(self, B: torch.Tensor, tp_rank: int):
+        return B
+
 
 def get_lora_layer(
     layer: nn.Module, lora_rank: int, scaling: int, lora_backend: BaseLoRABackend
diff --git a/python/sglang/srt/lora/lora.py b/python/sglang/srt/lora/lora.py
index ed9a48fc8fdb..51e4d56f2f8c 100644
--- a/python/sglang/srt/lora/lora.py
+++ b/python/sglang/srt/lora/lora.py
@@ -39,16 +39,9 @@ def __init__(self, config: LoRAConfig, base_hf_config: AutoConfig):
         super().__init__()
         self.config: LoRAConfig = config
         self.base_hf_config: AutoConfig = base_hf_config
-        self.weights: Dict[str, torch.Tensor] = {}
-        self.weight_gpu: Dict[str, torch.Tensor] = {}
-
-    def load_to_gpu(self):
-        for name, weight in self.weights.items():
-            self.weight_gpu[name] = weight.to(torch.float16).to("cuda")
 
-    def offload_from_gpu(self):
-        for name, weight in self.weights.items():
-            self.weight_gpu[name] = None
+        # lora weights in cpu. The weights are loaded from checkpoint.
+        self.weights: Dict[str, torch.Tensor] = {}
 
 
 class LoRAAdapter(nn.Module):
@@ -77,19 +70,6 @@ def __init__(
         )
 
         self.weights: Dict[str, torch.Tensor] = {}
-        self.weights_gpu: Dict[str, torch.Tensor] = {}
-
-    def load_to_gpu(self):
-        for name, weight in self.weights.items():
-            self.weights_gpu[name] = weight.to(torch.float16).to("cuda")
-        for layer in self.layers:
-            layer.load_to_gpu()
-
-    def offload_from_gpu(self):
-        for name, weight in self.weights.items():
-            self.weights_gpu[name] = None
-        for layer in self.layers:
-            layer.offload_from_gpu()
 
     # initialize the LoRA weights to cpu
     def initialize_weights(self):
diff --git a/python/sglang/srt/lora/lora_manager.py b/python/sglang/srt/lora/lora_manager.py
index 3396f0d8ec10..88f9072373cd 100644
--- a/python/sglang/srt/lora/lora_manager.py
+++ b/python/sglang/srt/lora/lora_manager.py
@@ -23,7 +23,7 @@
 from sglang.srt.configs.load_config import LoadConfig
 from sglang.srt.hf_transformers_utils import AutoConfig
 from sglang.srt.lora.backend import BaseLoRABackend, get_backend_from_name
-from sglang.srt.lora.layers import get_lora_layer
+from sglang.srt.lora.layers import BaseLayerWithLoRA, get_lora_layer
 from sglang.srt.lora.lora import LoRAAdapter
 from sglang.srt.lora.lora_config import LoRAConfig
 from sglang.srt.lora.mem_pool import LoRAMemoryPool
@@ -51,6 +51,8 @@ def __init__(
         load_config: LoadConfig,
         dtype: torch.dtype,
         lora_backend: str = "triton",
+        tp_size: int = 1,
+        tp_rank: int = 0,
     ):
         self.base_model: torch.nn.Module = base_model
         self.lora_paths: Dict[str, str] = lora_paths
@@ -58,6 +60,9 @@ def __init__(
         self.max_loras_per_batch: int = max_loras_per_batch
         self.load_config: LoadConfig = load_config
         self.dtype: torch.dtype = dtype
+        self.device: torch.device = next(self.base_model.parameters()).device
+        self.tp_size: int = tp_size
+        self.tp_rank: int = tp_rank
 
         # LoRA backend for running sgemm kernels
         logger.info(f"Using {lora_backend} as backend of LoRA kernels.")
@@ -110,7 +115,13 @@ def init_loras(self):
     def init_lora_memory_pool(self):
         # Initialize memory pool
         self.memory_pool = LoRAMemoryPool(
-            self.base_hf_config, self.max_loras_per_batch, self.max_lora_dim, self.dtype
+            self.base_hf_config,
+            self.max_loras_per_batch,
+            self.max_lora_dim,
+            self.dtype,
+            self.tp_size,
+            self.tp_rank,
+            self.lora_modules,
         )
 
         # Initialize target lora modules in memory pool
@@ -131,12 +142,12 @@ def prepare_lora_batch(self, forward_batch: ForwardBatch):
         seg_lens = (
             forward_batch.extend_seq_lens
             if forward_batch.forward_mode.is_extend()
-            else torch.ones(bs, device="cuda")
+            else torch.ones(bs, device=self.device)
         )
-        seg_indptr = torch.zeros((bs + 1,), dtype=torch.int32, device="cuda")
+        seg_indptr = torch.zeros((bs + 1,), dtype=torch.int32, device=self.device)
         seg_indptr[1:] = torch.cumsum(seg_lens, dim=0)
         max_len = int(torch.max(seg_lens))
-        weight_indices = torch.empty((bs,), dtype=torch.int64, device="cuda")
+        weight_indices = torch.empty((bs,), dtype=torch.int64, device=self.device)
         for i, lora_path in enumerate(forward_batch.lora_paths):
             weight_indices[i] = self.memory_pool.get_buffer_id(lora_path)
 
@@ -150,22 +161,32 @@ def prepare_lora_batch(self, forward_batch: ForwardBatch):
         self.lora_backend.set_batch_info(batch_info)
 
         # call set_lora_info for each lora modules
-        for module_name, module in self.lora_modules:
-            layer_id = get_layer_id(module_name)
-            if "qkv_proj" not in module_name:
-                weight_name = get_weight_name(
-                    module_name, self.lora_weight_names, LoRAType.LORA_A
-                )
-                module.set_lora_info(
-                    self.memory_pool.get_tensor(weight_name, layer_id, LoRAType.LORA_A),
-                    self.memory_pool.get_tensor(weight_name, layer_id, LoRAType.LORA_B),
-                )
-            else:
-                module.set_lora_info(
-                    self.memory_pool.get_tensor("qkv_proj", layer_id, LoRAType.LORA_A),
-                    self.memory_pool.get_tensor("q_proj", layer_id, LoRAType.LORA_B),
-                    self.memory_pool.get_tensor("kv_proj", layer_id, LoRAType.LORA_B),
-                )
+        for layer_id, modules in self.lora_modules.items():
+            for module_name, module in modules:
+                if "qkv_proj" in module_name:
+                    module.set_lora_info(
+                        self.memory_pool.get_tensor(
+                            "qkv_proj", layer_id, LoRAType.LORA_A
+                        ),
+                        self.memory_pool.get_tensor(
+                            "q_proj", layer_id, LoRAType.LORA_B
+                        ),
+                        self.memory_pool.get_tensor(
+                            "kv_proj", layer_id, LoRAType.LORA_B
+                        ),
+                    )
+                else:
+                    weight_name = get_weight_name(
+                        module_name, self.lora_weight_names, LoRAType.LORA_A
+                    )
+                    module.set_lora_info(
+                        self.memory_pool.get_tensor(
+                            weight_name, layer_id, LoRAType.LORA_A
+                        ),
+                        self.memory_pool.get_tensor(
+                            weight_name, layer_id, LoRAType.LORA_B
+                        ),
+                    )
 
     def set_lora_module(self, module_name, module):
         lora_module = get_lora_layer(
@@ -182,10 +203,13 @@ def convert_to_lora_layers(self):
         )
 
         # Monkey patch to use the LoRA version layers
-        self.lora_modules: List[Tuple[str, torch.nn.Module]] = []
+        self.lora_modules: Dict[int, List[Tuple[str, BaseLayerWithLoRA]]] = {
+            i: [] for i in range(self.base_hf_config.num_hidden_layers)
+        }
         for module_name, module in self.base_model.named_modules():
             # The module should be converted if it is included in target_names
             if module_name.split(".")[-1] in customized_target_names:
-                self.lora_modules.append(
+                layer_id = get_layer_id(module_name)
+                self.lora_modules[layer_id].append(
                     (module_name, self.set_lora_module(module_name, module))
                 )
diff --git a/python/sglang/srt/lora/mem_pool.py b/python/sglang/srt/lora/mem_pool.py
index a1a6e549dd8b..4c17a3c925e2 100644
--- a/python/sglang/srt/lora/mem_pool.py
+++ b/python/sglang/srt/lora/mem_pool.py
@@ -2,9 +2,12 @@
 
 import torch
 
+from sglang.srt.distributed import divide
 from sglang.srt.hf_transformers_utils import AutoConfig
+from sglang.srt.lora.layers import BaseLayerWithLoRA
 from sglang.srt.lora.lora import LoRAAdapter
 from sglang.srt.lora.utils import (
+    ROW_PARALLELISM_LINEAR_LORA_NAMES,
     LoRAType,
     get_hidden_dim,
     get_stacked_multiply,
@@ -21,6 +24,9 @@ def __init__(
         max_loras_per_batch: int,
         max_lora_dim: int,
         dtype: torch.dtype,
+        tp_size: int,
+        tp_rank: int,
+        lora_modules: Dict[int, List[Tuple[str, BaseLayerWithLoRA]]],
     ):
 
         self.base_hf_config: AutoConfig = base_hf_config
@@ -28,6 +34,9 @@ def __init__(
         self.max_loras_per_batch: int = max_loras_per_batch
         self.max_lora_dim: int = max_lora_dim
         self.dtype: torch.dtype = dtype
+        self.tp_size: int = tp_size
+        self.tp_rank: int = tp_rank
+        self.lora_modules: Dict[int, List[Tuple[str, BaseLayerWithLoRA]]] = lora_modules
 
         # Both A_buffer and B_buffer maps lora weight names to its buffer space.
         # A_buffer contains num_layer number of row-major tensors with shape
@@ -45,6 +54,41 @@ def __init__(
         # Here we don't initalize to None since None is a valid uid
         self.buffer_id_to_uid: List[Optional[str]] = [""] * self.max_loras_per_batch
 
+    def get_lora_A_shape(
+        self, module_name: str, base_model: torch.nn.Module
+    ) -> Tuple[int]:
+        """
+        Given a module_name (might be a stacked name), return the hidden dims of modules's input and output.
+        """
+        input_dim, _ = get_hidden_dim(module_name, self.base_hf_config, base_model)
+        c = get_stacked_multiply(module_name)
+        if self.tp_size > 1:
+            if module_name in ROW_PARALLELISM_LINEAR_LORA_NAMES:
+                input_dim = divide(input_dim, self.tp_size)
+        return (
+            self.max_loras_per_batch,
+            self.max_lora_dim * c,
+            input_dim,
+        )
+
+    def get_lora_B_shape(
+        self, module_name: str, base_model: torch.nn.Module
+    ) -> Tuple[int]:
+        """
+        Given a module_name (might be a stacked name), return the hidden dims of modules's input and output.
+        """
+        _, output_dim = get_hidden_dim(module_name, self.base_hf_config, base_model)
+        c = get_stacked_multiply(module_name)
+        if self.tp_size > 1:
+            if module_name not in ROW_PARALLELISM_LINEAR_LORA_NAMES:
+                output_dim = divide(output_dim, self.tp_size)
+        return (
+            c,
+            self.max_loras_per_batch,
+            output_dim,
+            self.max_lora_dim,
+        )
+
     def init_buffers(
         self,
         lora_weight_names: Set[Tuple[str]],
@@ -54,42 +98,31 @@ def init_buffers(
         # lora_weight_names is a set of name pairs indicating each pair of lora modules to load
         #   e.g., {("qkv_proj", "q_proj"), ("qkv_proj", "kv_proj"), ("o_proj", "o_proj")}
         self.lora_weight_names: Set[Tuple[str]] = lora_weight_names
-
-        for module_A, module_B in lora_weight_names:
-            # Init A tensor, column_major=False
-            input_dim, _ = get_hidden_dim(module_A, self.base_hf_config, base_model)
-            c = get_stacked_multiply(module_A)
-            if module_A not in self.A_buffer:
-                self.A_buffer[module_A] = [
-                    torch.empty(
-                        (
-                            self.max_loras_per_batch,
-                            self.max_lora_dim * c,
-                            input_dim,
-                        ),
-                        dtype=self.dtype,
-                        device="cuda",
-                    )
-                    for i in range(self.num_layer)
-                ]
-
-            # Init B tensor, column_major=True
-            _, output_dim = get_hidden_dim(module_B, self.base_hf_config, base_model)
-            c = get_stacked_multiply(module_B)
-            if module_B not in self.B_buffer:
-                self.B_buffer[module_B] = [
-                    torch.empty(
-                        (
-                            c,  # stacked lora_b modules might need separation
-                            self.max_loras_per_batch,
-                            output_dim,
-                            self.max_lora_dim,
-                        ),
-                        dtype=self.dtype,
-                        device="cuda",
-                    )
-                    for i in range(self.num_layer)
-                ]
+        device = next(base_model.parameters()).device
+        lora_module_A_names = set([name[0] for name in lora_weight_names])
+        lora_module_B_names = set([name[1] for name in lora_weight_names])
+        # Init A tensor, column_major=False
+        for module_A in lora_module_A_names:
+            lora_A_shape = self.get_lora_A_shape(module_A, base_model)
+            self.A_buffer[module_A] = [
+                torch.empty(
+                    lora_A_shape,
+                    dtype=self.dtype,
+                    device=device,
+                )
+                for i in range(self.num_layer)
+            ]
+        # Init B tensor, column_major=True
+        for module_B in lora_module_B_names:
+            lora_B_shape = self.get_lora_B_shape(module_B, base_model)
+            self.B_buffer[module_B] = [
+                torch.empty(
+                    lora_B_shape,
+                    dtype=self.dtype,
+                    device=device,
+                )
+                for _ in range(self.num_layer)
+            ]
 
     def prepare_lora_batch(
         self,
@@ -136,30 +169,56 @@ def load_lora_weight_to_buffer(
         assert lora_adapter is not None
         for layer_id in range(self.num_layer):
             layer_weights = lora_adapter.layers[layer_id].weights
+            temp_A_buffer: Dict[str, torch.Tensor] = {}
+            temp_B_buffer: Dict[str, torch.Tensor] = {}
             for name, weights in layer_weights.items():
                 if "lora_A" in name:
                     lora_weight_name = get_weight_name(
                         name, self.lora_weight_names, LoRAType.LORA_A
                     )
-                    if lora_weight_name:
-                        self.A_buffer[lora_weight_name][layer_id][buffer_id].copy_(
-                            weights
-                        )
+                    temp_A_buffer[lora_weight_name] = weights
                 else:
                     lora_weight_name = get_weight_name(
                         name, self.lora_weight_names, LoRAType.LORA_B
                     )
-                    if lora_weight_name:
-                        c = get_stacked_multiply(lora_weight_name)
-                        if c > 1:
-                            for stacked_id in range(c):
-                                self.B_buffer[lora_weight_name][layer_id][stacked_id][
-                                    buffer_id
-                                ].copy_(weights[stacked_id])
-                        else:
-                            self.B_buffer[lora_weight_name][layer_id][0][
-                                buffer_id
-                            ].copy_(weights)
+                    temp_B_buffer[lora_weight_name] = weights
+
+            if self.tp_size > 1:
+                cur_layer_modules = self.lora_modules[layer_id]
+                for module_name, module in cur_layer_modules:
+                    if "qkv_proj" in module_name:
+                        temp_A_buffer["qkv_proj"] = module.slice_lora_a_weights(
+                            temp_A_buffer["qkv_proj"], self.tp_rank
+                        )
+                        temp_B_buffer["q_proj"], temp_B_buffer["kv_proj"] = (
+                            module.slice_lora_b_weights(
+                                [temp_B_buffer["q_proj"], temp_B_buffer["kv_proj"]],
+                                self.tp_rank,
+                            )
+                        )
+                    else:
+                        weight_name = get_weight_name(
+                            module_name, self.lora_weight_names, LoRAType.LORA_A
+                        )
+                        temp_A_buffer[weight_name] = module.slice_lora_a_weights(
+                            temp_A_buffer[weight_name], self.tp_rank
+                        )
+                        temp_B_buffer[weight_name] = module.slice_lora_b_weights(
+                            temp_B_buffer[weight_name], self.tp_rank
+                        )
+
+            for name, weights in temp_A_buffer.items():
+                self.A_buffer[name][layer_id][buffer_id].copy_(weights)
+
+            for name, weights in temp_B_buffer.items():
+                c = get_stacked_multiply(name)
+                if c > 1:
+                    for stacked_id in range(c):
+                        self.B_buffer[name][layer_id][stacked_id][buffer_id].copy_(
+                            weights[stacked_id]
+                        )
+                else:
+                    self.B_buffer[name][layer_id][0][buffer_id].copy_(weights)
 
     def get_tensor(
         self, weight_name: str, layer_id: int, lora_type: LoRAType
diff --git a/python/sglang/srt/lora/utils.py b/python/sglang/srt/lora/utils.py
index f5845373d338..c53336ba6aa0 100644
--- a/python/sglang/srt/lora/utils.py
+++ b/python/sglang/srt/lora/utils.py
@@ -133,9 +133,20 @@ def get_weight_name(
     target_name is name of a given module,
     lora_weight_names is a set of lora stacked name pairs (see get_stacked_name method above)
     If there is a weight name in lora_weight_names that can match target_name, return this name
-    Else return None
+    Else raise ValueError.
     """
     idx = 0 if lora_type == LoRAType.LORA_A else 1
     for weight_name_pair in lora_weight_names:
         if weight_name_pair[idx] in target_name:
             return weight_name_pair[idx]
+    raise ValueError(
+        f"Cannot find weight name for {target_name} in {lora_weight_names}"
+    )
+
+
+# TODO: [PR #4274] For future use to simplify the mapping between HF module names and customized module names.
+VOCAB_PARALLELISM_EMBEDDING_NAMES = ["embeddings"]
+COLUMN_PARALLELISM_LINEAR_LORA_NAMES = ["gate_proj", "up_proj"]
+MERGED_COLUMN_PARALLELISM_LINEAR_LORA_NAMES = ["gate_up_proj"]
+QKV_PARALLELISM_LINEAR_LORA_NAMES = ["qkv_proj"]
+ROW_PARALLELISM_LINEAR_LORA_NAMES = ["o_proj", "down_proj"]
diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py
index 55558e0f62aa..6f2e05e269d9 100644
--- a/python/sglang/srt/model_executor/model_runner.py
+++ b/python/sglang/srt/model_executor/model_runner.py
@@ -188,9 +188,6 @@ def initialize(self, min_per_gpu_memory: float):
         supports_torch_tp = getattr(self.model, "supports_torch_tp", False)
         if self.tp_size > 1 and supports_torch_tp:
             self.apply_torch_tp()
-            self.torch_tp_applied = True
-        else:
-            self.torch_tp_applied = False
 
         # Init lora
         if server_args.lora_paths is not None:
@@ -624,6 +621,8 @@ def init_lora_manager(self):
             load_config=self.load_config,
             dtype=self.dtype,
             lora_backend=self.server_args.lora_backend,
+            tp_size=self.tp_size,
+            tp_rank=self.tp_rank,
         )
         logger.info("LoRA manager ready.")
 
diff --git a/python/sglang/srt/utils.py b/python/sglang/srt/utils.py
index 8a8be455006b..7078016040c9 100644
--- a/python/sglang/srt/utils.py
+++ b/python/sglang/srt/utils.py
@@ -257,7 +257,7 @@ def get_available_gpu_memory(device, gpu_id, distributed=False, empty_cache=True
     When distributed is True, the available memory is the minimum available memory of all GPUs.
     """
     if device == "cuda":
-        num_gpus = torch.cuda.device_count()
+        num_gpus = cuda_device_count_stateless()
         assert gpu_id < num_gpus
 
         if torch.cuda.current_device() != gpu_id:
diff --git a/python/sglang/test/runners.py b/python/sglang/test/runners.py
index d92dca4f0783..17b5cc812c06 100644
--- a/python/sglang/test/runners.py
+++ b/python/sglang/test/runners.py
@@ -437,6 +437,7 @@ def __init__(
         speculative_eagle_topk: Optional[int] = None,
         speculative_num_draft_tokens: Optional[int] = None,
         disable_overlap_schedule: bool = False,
+        disable_custom_all_reduce: bool = False,
     ):
         self.model_type = model_type
         self.is_generation = model_type == "generation"
@@ -470,6 +471,7 @@ def __init__(
             enable_ep_moe=enable_ep_moe,
             disable_overlap_schedule=disable_overlap_schedule,
             cuda_graph_max_bs=4,
+            disable_custom_all_reduce=disable_custom_all_reduce,
             **spec_kwargs,
         )
 
diff --git a/test/srt/models/lora/test_lora_backend.py b/test/srt/models/lora/test_lora_backend.py
index 08d3494a60bf..4c7fe80748b9 100644
--- a/test/srt/models/lora/test_lora_backend.py
+++ b/test/srt/models/lora/test_lora_backend.py
@@ -49,7 +49,7 @@
     LoRAModelCase(
         base="meta-llama/Llama-2-7b-hf",
         adaptors=[LoRAAdaptor(name="winddude/wizardLM-LlaMA-LoRA-7B")],
-        max_loras_per_batch=1,
+        max_loras_per_batch=2,
     ),
 ]
 
@@ -96,6 +96,7 @@ def run_backend(
             disable_cuda_graph=True,
             disable_radix_cache=True,
             mem_fraction_static=0.88,
+            disable_custom_all_reduce=False,
         ) as srt_runner:
             srt_outputs = srt_runner.forward(
                 [prompt], max_new_tokens=max_new_tokens, lora_paths=[adaptor.name]
@@ -114,6 +115,7 @@ def run_backend(
             model_type="generation",
             tp_size=model_case.tp_size,
             mem_fraction_static=0.88,
+            disable_custom_all_reduce=False,
         ) as srt_runner:
             srt_no_lora_outputs = srt_runner.forward(
                 [prompt], max_new_tokens=max_new_tokens
diff --git a/test/srt/models/lora/test_lora_tp.py b/test/srt/models/lora/test_lora_tp.py
new file mode 100644
index 000000000000..2da0d9e02bc5
--- /dev/null
+++ b/test/srt/models/lora/test_lora_tp.py
@@ -0,0 +1,257 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import multiprocessing as mp
+import os
+import unittest
+from typing import List
+
+import torch
+from utils import TORCH_DTYPES, LoRAAdaptor, LoRAModelCase
+
+from sglang.test.runners import HFRunner, SRTRunner
+from sglang.test.test_utils import calculate_rouge_l, is_in_ci
+
+CI_LORA_MODELS = [
+    LoRAModelCase(
+        base="meta-llama/Llama-3.1-8B-Instruct",
+        adaptors=[
+            LoRAAdaptor(
+                name="algoprog/fact-generation-llama-3.1-8b-instruct-lora",
+            ),
+        ],
+        max_loras_per_batch=1,
+    ),
+    LoRAModelCase(
+        base="meta-llama/Llama-3.1-8B-Instruct",
+        adaptors=[
+            LoRAAdaptor(
+                name="Nutanix/Meta-Llama-3.1-8B-Instruct_lora_4_alpha_16",
+                prefill_tolerance=1e-1,
+            ),
+        ],
+        max_loras_per_batch=1,
+    ),
+]
+
+ALL_OTHER_LORA_MODELS = [
+    LoRAModelCase(
+        base="meta-llama/Llama-2-7b-hf",
+        adaptors=[LoRAAdaptor(name="winddude/wizardLM-LlaMA-LoRA-7B")],
+        max_loras_per_batch=2,
+    ),
+]
+
+PROMPTS = [
+    "AI is a field of computer science focused on",
+    """
+    ### Instruction:
+    Tell me about llamas and alpacas
+    ### Response:
+    Llamas are large, long-necked animals with a woolly coat. They have two toes on each foot instead of three like other camelids (camels, dromedaries). Llamas live in the Andean mountains of South America where they graze on grasses and shrubs. Alpaca is another name for domesticated llama. The word "alpaca" comes from an Incan language meaning "golden fleece." Alpacas look very similar to llamas but are smaller than their wild relatives. Both species were used by ancient people as pack animals and for meat. Today both llamas and alpacas are raised primarily for their fiber which can be spun into yarn or knitted into clothing.
+    ### Question 2:
+    What do you know about llamas?
+    ### Answer:
+    """,
+]
+
+BACKEND = "triton"
+
+
+class TestLoRATP(unittest.TestCase):
+    def run_tp(
+        self,
+        prompt: str,
+        model_case: LoRAModelCase,
+        torch_dtype: torch.dtype,
+        max_new_tokens: int,
+    ):
+        """
+        Run triton backend tests with specified TP size for a single prompt and model case.
+        """
+        base_path = model_case.base
+        adaptor = model_case.adaptors[0]
+        tp_size = model_case.tp_size
+        print(
+            f"\n========== Testing triton backend with TP size {tp_size} for base '{base_path}' --- "
+            f"Prompt '{prompt[:50]}...' using adaptor '{adaptor.name}' ---"
+        )
+        with SRTRunner(
+            base_path,
+            torch_dtype=torch_dtype,
+            model_type="generation",
+            tp_size=tp_size,
+            lora_paths=[adaptor.name for adaptor in model_case.adaptors],
+            max_loras_per_batch=model_case.max_loras_per_batch,
+            lora_backend=BACKEND,
+            disable_cuda_graph=True,
+            disable_radix_cache=True,
+            mem_fraction_static=0.88,
+            disable_custom_all_reduce=True,
+        ) as srt_runner:
+            srt_outputs = srt_runner.forward(
+                [prompt], max_new_tokens=max_new_tokens, lora_paths=[adaptor.name]
+            )
+
+        with HFRunner(
+            base_path, torch_dtype=torch_dtype, model_type="generation"
+        ) as hf_runner:
+            hf_outputs = hf_runner.forward(
+                [prompt], max_new_tokens=max_new_tokens, lora_paths=[adaptor.name]
+            )
+
+        with SRTRunner(
+            base_path,
+            torch_dtype=torch_dtype,
+            model_type="generation",
+            tp_size=tp_size,
+            mem_fraction_static=0.88,
+            disable_custom_all_reduce=True,
+        ) as srt_runner:
+            srt_no_lora_outputs = srt_runner.forward(
+                [prompt], max_new_tokens=max_new_tokens
+            )
+
+        with HFRunner(
+            base_path,
+            torch_dtype=torch_dtype,
+            model_type="generation",
+        ) as hf_runner:
+            hf_no_lora_outputs = hf_runner.forward(
+                [prompt], max_new_tokens=max_new_tokens
+            )
+
+        # Use individual adapter tolerances if set, otherwise use model defaults
+        prefill_tol = (
+            adaptor.prefill_tolerance
+            if adaptor.prefill_tolerance is not None
+            else model_case.prefill_tolerance
+        )
+        decode_tol = (
+            adaptor.decode_tolerance
+            if adaptor.decode_tolerance is not None
+            else model_case.decode_tolerance
+        )
+        rouge_tol = (
+            adaptor.rouge_l_tolerance
+            if adaptor.rouge_l_tolerance is not None
+            else model_case.rouge_l_tolerance
+        )
+
+        # Compare prefill stage logprobs (HF vs SRTRunner with LoRA)
+        hf_prefill = torch.tensor(hf_outputs.top_input_logprobs[0])
+        srt_prefill = torch.tensor(srt_outputs.top_input_logprobs[0])
+        max_prefill_diff = torch.max(torch.abs(hf_prefill - srt_prefill))
+        print("Max prefill diff (HF vs SRT):", max_prefill_diff)
+
+        # Compare decode stage logprobs
+        hf_decode = torch.tensor(hf_outputs.top_output_logprobs[0])
+        srt_decode = torch.tensor(srt_outputs.top_output_logprobs[0])
+        max_decode_diff = torch.max(torch.abs(hf_decode - srt_decode))
+        print("Max decode diff (HF vs SRT):", max_decode_diff)
+
+        srt_output_str = srt_outputs.output_strs[0].strip()
+        hf_output_str = hf_outputs.output_strs[0].strip()
+        rouge_score = calculate_rouge_l([srt_output_str], [hf_output_str])[0]
+        print("ROUGE-L score:", rouge_score)
+        print("SRT output:", srt_output_str)
+        print("HF output:", hf_output_str)
+
+        # Additional: compare prefill outputs between base model (no LoRA) and LoRA model for reference
+        hf_no_lora_prefill = torch.tensor(hf_no_lora_outputs.top_input_logprobs[0])
+        srt_no_lora_prefill = torch.tensor(srt_no_lora_outputs.top_input_logprobs[0])
+        print(
+            "Max diff (SRT base vs SRT LoRA prefill):",
+            torch.max(torch.abs(srt_no_lora_prefill - srt_prefill)),
+        )
+        print(
+            "Max diff (HF base vs HF LoRA prefill):",
+            torch.max(torch.abs(hf_no_lora_prefill - hf_prefill)),
+        )
+
+        if hf_prefill.shape[0] <= 100:
+            assert torch.all(torch.abs(hf_prefill - srt_prefill) < prefill_tol), (
+                f"Prefill logprobs mismatch for base '{base_path}', adaptor '{adaptor.name}', "
+                f"triton backend with TP {tp_size}, prompt: '{prompt[:50]}...'"
+            )
+
+        if hf_decode.shape[0] <= 100:
+            assert torch.all(torch.abs(hf_decode - srt_decode) < decode_tol), (
+                f"Decode logprobs mismatch for base '{base_path}', adaptor '{adaptor.name}', "
+                f"triton backend with TP {tp_size}, prompt: '{prompt[:50]}...'"
+            )
+
+        if rouge_score < rouge_tol:
+            raise AssertionError(
+                f"ROUGE-L score {rouge_score} below tolerance {rouge_tol} "
+                f"for base '{base_path}', adaptor '{adaptor.name}', triton backend with TP {tp_size}, prompt: '{prompt[:50]}...'"
+            )
+
+    def run_tp_batch(
+        self,
+        prompts: List[str],
+        model_case: LoRAModelCase,
+        torch_dtype: torch.dtype,
+        max_new_tokens: int,
+        tp_size: int,
+    ):
+        # TODO: Implement batch processing version of run_tp
+        raise NotImplementedError(
+            "Batch processing version of run_tp is not implemented yet."
+        )
+
+    def _run_tp_on_model_cases(self, model_cases: List[LoRAModelCase]):
+        tp_list = [2]  # Define TP sizes to iterate over
+        for model_case in model_cases:
+            # If skip_long_prompt is True, filter out prompts longer than 1000 characters
+            prompts = (
+                PROMPTS
+                if not model_case.skip_long_prompt
+                else [p for p in PROMPTS if len(p) < 1000]
+            )
+            for tp_size in tp_list:
+                model_case.tp_size = tp_size
+                for torch_dtype in TORCH_DTYPES:
+                    for prompt in prompts:
+                        self.run_tp(
+                            prompt,
+                            model_case,
+                            torch_dtype,
+                            max_new_tokens=32,
+                        )
+
+    def test_ci_lora_models(self):
+        self._run_tp_on_model_cases(CI_LORA_MODELS)
+
+    def test_all_lora_models(self):
+        if is_in_ci():
+            return
+
+        # Retain ONLY_RUN check here
+        filtered_models = []
+        for model_case in ALL_OTHER_LORA_MODELS:
+            if "ONLY_RUN" in os.environ and os.environ["ONLY_RUN"] != model_case.base:
+                continue
+            filtered_models.append(model_case)
+
+        self._run_tp_on_model_cases(filtered_models)
+
+
+if __name__ == "__main__":
+    try:
+        mp.set_start_method("spawn")
+    except RuntimeError:
+        pass
+
+    unittest.main(warnings="ignore")


Reviews:
user: login: Fridge003
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-11T01:05:25Z
user: login: aoshen524
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-11T01:33:55Z
user: login: aoshen524
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-11T01:34:43Z
user: login: aoshen524
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-11T01:35:16Z
user: login: aoshen524
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-11T01:46:33Z
user: login: Fridge003
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-11T17:46:34Z
user: login: aoshen524
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-11T18:30:47Z
user: login: Fridge003
body: LGTM
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-03-15T18:17:46Z
user: login: Ying1123
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-03-19T03:05:38Z

Review comments:
diff_hunk: @@ -142,6 +169,16 @@ def apply_lora(self, base_output: torch.Tensor, x: torch.Tensor) -> torch.Tensor
             else base_output + lora_output * self.scaling
         )
 
+    def slice_lora_a_weights(self, A: torch.Tensor, tp_rank: int):
+        return A
+
+    def slice_lora_b_weights(self, B: torch.Tensor, tp_rank: int):
+        # 是一样的，所以取index 0
path: python/sglang/srt/lora/layers.py
user: login: Fridge003
body: Please remove the Chinese comment here
created_at: 2025-03-11T00:23:15Z
author_association: COLLABORATOR
diff_hunk: @@ -133,9 +133,15 @@ def get_weight_name(
     target_name is name of a given module,
     lora_weight_names is a set of lora stacked name pairs (see get_stacked_name method above)
     If there is a weight name in lora_weight_names that can match target_name, return this name
-    Else return None
+    Else raise ValueError.
     """
     idx = 0 if lora_type == LoRAType.LORA_A else 1
     for weight_name_pair in lora_weight_names:
         if weight_name_pair[idx] in target_name:
             return weight_name_pair[idx]
+    raise ValueError(
+        f"Cannot find weight name for {target_name} in {lora_weight_names}"
+    )
+
+
+ROW_PARALLELISM_LORA_NAMES = ["o_proj", "down_proj"]
path: python/sglang/srt/lora/utils.py
user: login: Fridge003
body: I feel `COL_PARALLELISM_LORA_NAMES` can also be added, maybe it will be useful in the future
created_at: 2025-03-11T00:45:27Z
author_association: COLLABORATOR
diff_hunk: @@ -175,9 +175,6 @@ def __init__(
         supports_torch_tp = getattr(self.model, "supports_torch_tp", False)
         if self.tp_size > 1 and supports_torch_tp:
             self.apply_torch_tp()
-            self.torch_tp_applied = True
path: python/sglang/srt/model_executor/model_runner.py
user: login: Fridge003
body: Why removing logic about `self.torch_tp_applied` here
created_at: 2025-03-11T00:46:30Z
author_association: COLLABORATOR
diff_hunk: @@ -42,14 +43,16 @@
             ),
         ],
         max_loras_per_batch=1,
+        tp_size=2,
     ),
 ]
 
 ALL_OTHER_LORA_MODELS = [
     LoRAModelCase(
         base="meta-llama/Llama-2-7b-hf",
         adaptors=[LoRAAdaptor(name="winddude/wizardLM-LlaMA-LoRA-7B")],
-        max_loras_per_batch=1,
+        max_loras_per_batch=2,
+        tp_size=2,
     ),
 ]
path: test/srt/models/lora/test_lora_backend.py
user: login: Fridge003
body: You can add a new set of test (even a new test file) for tensor parallelism, which uses its own set of `TP_LORA_MODELS` as model cases. The original tests with `tp_size=1` should be kept.
created_at: 2025-03-11T00:49:14Z
author_association: COLLABORATOR
diff_hunk: @@ -21,13 +24,20 @@ def __init__(
         max_loras_per_batch: int,
         max_lora_dim: int,
         dtype: torch.dtype,
+        tp_size: int,
+        tp_rank: int,
+        lora_modules: List[Tuple[str, BaseLayerWithLoRA]],
     ):
 
         self.base_hf_config: AutoConfig = base_hf_config
         self.num_layer: int = base_hf_config.num_hidden_layers
         self.max_loras_per_batch: int = max_loras_per_batch
         self.max_lora_dim: int = max_lora_dim
         self.dtype: torch.dtype = dtype
+        self.tp_size: int = tp_size
+        self.tp_rank: int = tp_rank
+        self.lora_modules: List[Tuple[str, BaseLayerWithLoRA]] = lora_modules
+        self.num_modules_per_layer = divide(len(self.lora_modules), self.num_layer)
path: python/sglang/srt/lora/mem_pool.py
user: login: Fridge003
body: This will cause bug after we add embedding and lm_head layers. Is there other way to get `self.num_modules_per_layer`
created_at: 2025-03-11T00:56:56Z
author_association: COLLABORATOR
diff_hunk: @@ -142,6 +169,16 @@ def apply_lora(self, base_output: torch.Tensor, x: torch.Tensor) -> torch.Tensor
             else base_output + lora_output * self.scaling
         )
 
+    def slice_lora_a_weights(self, A: torch.Tensor, tp_rank: int):
+        return A
+
+    def slice_lora_b_weights(self, B: torch.Tensor, tp_rank: int):
+        # 是一样的，所以取index 0
path: python/sglang/srt/lora/layers.py
user: login: aoshen524
body: sure
created_at: 2025-03-11T01:33:55Z
author_association: CONTRIBUTOR
diff_hunk: @@ -42,14 +43,16 @@
             ),
         ],
         max_loras_per_batch=1,
+        tp_size=2,
     ),
 ]
 
 ALL_OTHER_LORA_MODELS = [
     LoRAModelCase(
         base="meta-llama/Llama-2-7b-hf",
         adaptors=[LoRAAdaptor(name="winddude/wizardLM-LlaMA-LoRA-7B")],
-        max_loras_per_batch=1,
+        max_loras_per_batch=2,
+        tp_size=2,
     ),
 ]
path: test/srt/models/lora/test_lora_backend.py
user: login: aoshen524
body: sure

created_at: 2025-03-11T01:34:43Z
author_association: CONTRIBUTOR
diff_hunk: @@ -21,13 +24,20 @@ def __init__(
         max_loras_per_batch: int,
         max_lora_dim: int,
         dtype: torch.dtype,
+        tp_size: int,
+        tp_rank: int,
+        lora_modules: List[Tuple[str, BaseLayerWithLoRA]],
     ):
 
         self.base_hf_config: AutoConfig = base_hf_config
         self.num_layer: int = base_hf_config.num_hidden_layers
         self.max_loras_per_batch: int = max_loras_per_batch
         self.max_lora_dim: int = max_lora_dim
         self.dtype: torch.dtype = dtype
+        self.tp_size: int = tp_size
+        self.tp_rank: int = tp_rank
+        self.lora_modules: List[Tuple[str, BaseLayerWithLoRA]] = lora_modules
+        self.num_modules_per_layer = divide(len(self.lora_modules), self.num_layer)
path: python/sglang/srt/lora/mem_pool.py
user: login: aoshen524
body: nice insight!

created_at: 2025-03-11T01:35:16Z
author_association: CONTRIBUTOR
diff_hunk: @@ -133,9 +133,15 @@ def get_weight_name(
     target_name is name of a given module,
     lora_weight_names is a set of lora stacked name pairs (see get_stacked_name method above)
     If there is a weight name in lora_weight_names that can match target_name, return this name
-    Else return None
+    Else raise ValueError.
     """
     idx = 0 if lora_type == LoRAType.LORA_A else 1
     for weight_name_pair in lora_weight_names:
         if weight_name_pair[idx] in target_name:
             return weight_name_pair[idx]
+    raise ValueError(
+        f"Cannot find weight name for {target_name} in {lora_weight_names}"
+    )
+
+
+ROW_PARALLELISM_LORA_NAMES = ["o_proj", "down_proj"]
path: python/sglang/srt/lora/utils.py
user: login: aoshen524
body: sure

created_at: 2025-03-11T01:46:33Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,258 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import multiprocessing as mp
+import os
+import unittest
+from typing import List
+
+import torch
+from utils import BACKENDS, TORCH_DTYPES, LoRAAdaptor, LoRAModelCase
+
+from sglang.test.runners import HFRunner, SRTRunner
+from sglang.test.test_utils import calculate_rouge_l, is_in_ci
+
+CI_LORA_MODELS = [
+    LoRAModelCase(
+        base="meta-llama/Llama-3.1-8B-Instruct",
+        adaptors=[
+            LoRAAdaptor(
+                name="algoprog/fact-generation-llama-3.1-8b-instruct-lora",
+            ),
+        ],
+        max_loras_per_batch=1,
+        tp_size=2,
+    ),
+    LoRAModelCase(
+        base="meta-llama/Llama-3.1-8B-Instruct",
+        adaptors=[
+            LoRAAdaptor(
+                name="Nutanix/Meta-Llama-3.1-8B-Instruct_lora_4_alpha_16",
+                prefill_tolerance=1e-1,
+            ),
+        ],
+        max_loras_per_batch=1,
+        tp_size=2,
+    ),
+]
+
+ALL_OTHER_LORA_MODELS = [
+    LoRAModelCase(
+        base="meta-llama/Llama-2-7b-hf",
+        adaptors=[LoRAAdaptor(name="winddude/wizardLM-LlaMA-LoRA-7B")],
+        max_loras_per_batch=2,
+        tp_size=2,
+    ),
+]
+
+PROMPTS = [
+    "AI is a field of computer science focused on",
+    """
+    ### Instruction:
+    Tell me about llamas and alpacas
+    ### Response:
+    Llamas are large, long-necked animals with a woolly coat. They have two toes on each foot instead of three like other camelids (camels, dromedaries). Llamas live in the Andean mountains of South America where they graze on grasses and shrubs. Alpaca is another name for domesticated llama. The word "alpaca" comes from an Incan language meaning "golden fleece." Alpacas look very similar to llamas but are smaller than their wild relatives. Both species were used by ancient people as pack animals and for meat. Today both llamas and alpacas are raised primarily for their fiber which can be spun into yarn or knitted into clothing.
+    ### Question 2:
+    What do you know about llamas?
+    ### Answer:
+    """,
+]
+
+
+class TestLoRABackend(unittest.TestCase):
+    def run_backend(
+        self,
+        prompt: str,
+        model_case: LoRAModelCase,
+        torch_dtype: torch.dtype,
+        max_new_tokens: int,
+        backend: str,
+    ):
+        """
+        Run backend tests for a single prompt and model case.
+        """
+        base_path = model_case.base
+        adaptor = model_case.adaptors[0]
+        print(
+            f"\n========== Testing backend '{backend}' for base '{base_path}' --- "
+            f"Prompt '{prompt[:50]}...' using adaptor '{adaptor.name}' ---"
+        )
+        with SRTRunner(
+            base_path,
+            torch_dtype=torch_dtype,
+            model_type="generation",
+            tp_size=model_case.tp_size,
+            lora_paths=[adaptor.name for adaptor in model_case.adaptors],
+            max_loras_per_batch=model_case.max_loras_per_batch,
+            lora_backend=backend,
+            disable_cuda_graph=True,
+            disable_radix_cache=True,
+            mem_fraction_static=0.88,
+            disable_custom_all_reduce=True,
+        ) as srt_runner:
+            srt_outputs = srt_runner.forward(
+                [prompt], max_new_tokens=max_new_tokens, lora_paths=[adaptor.name]
+            )
+
+        with HFRunner(
+            base_path, torch_dtype=torch_dtype, model_type="generation"
+        ) as hf_runner:
+            hf_outputs = hf_runner.forward(
+                [prompt], max_new_tokens=max_new_tokens, lora_paths=[adaptor.name]
+            )
+
+        with SRTRunner(
+            base_path,
+            torch_dtype=torch_dtype,
+            model_type="generation",
+            tp_size=model_case.tp_size,
+            mem_fraction_static=0.88,
+            disable_custom_all_reduce=True,
+        ) as srt_runner:
+            srt_no_lora_outputs = srt_runner.forward(
+                [prompt], max_new_tokens=max_new_tokens
+            )
+
+        with HFRunner(
+            base_path,
+            torch_dtype=torch_dtype,
+            model_type="generation",
+        ) as hf_runner:
+            hf_no_lora_outputs = hf_runner.forward(
+                [prompt], max_new_tokens=max_new_tokens
+            )
+
+        # Use individual adapter tolerances if set, otherwise use model defaults
+        prefill_tol = (
+            adaptor.prefill_tolerance
+            if adaptor.prefill_tolerance is not None
+            else model_case.prefill_tolerance
+        )
+        decode_tol = (
+            adaptor.decode_tolerance
+            if adaptor.decode_tolerance is not None
+            else model_case.decode_tolerance
+        )
+        rouge_tol = (
+            adaptor.rouge_l_tolerance
+            if adaptor.rouge_l_tolerance is not None
+            else model_case.rouge_l_tolerance
+        )
+
+        # Compare prefill stage logprobs (HF vs SRTRunner with LoRA)
+        hf_prefill = torch.tensor(hf_outputs.top_input_logprobs[0])
+        srt_prefill = torch.tensor(srt_outputs.top_input_logprobs[0])
+        max_prefill_diff = torch.max(torch.abs(hf_prefill - srt_prefill))
+        print("Max prefill diff (HF vs SRT):", max_prefill_diff)
+
+        # Compare decode stage logprobs
+        hf_decode = torch.tensor(hf_outputs.top_output_logprobs[0])
+        srt_decode = torch.tensor(srt_outputs.top_output_logprobs[0])
+        max_decode_diff = torch.max(torch.abs(hf_decode - srt_decode))
+        print("Max decode diff (HF vs SRT):", max_decode_diff)
+
+        srt_output_str = srt_outputs.output_strs[0].strip()
+        hf_output_str = hf_outputs.output_strs[0].strip()
+        rouge_score = calculate_rouge_l([srt_output_str], [hf_output_str])[0]
+        print("ROUGE-L score:", rouge_score)
+        print("SRT output:", srt_output_str)
+        print("HF output:", hf_output_str)
+
+        # Additional: compare prefill outputs between base model (no LoRA) and LoRA model for reference
+        hf_no_lora_prefill = torch.tensor(hf_no_lora_outputs.top_input_logprobs[0])
+        srt_no_lora_prefill = torch.tensor(srt_no_lora_outputs.top_input_logprobs[0])
+        print(
+            "Max diff (SRT base vs SRT LoRA prefill):",
+            torch.max(torch.abs(srt_no_lora_prefill - srt_prefill)),
+        )
+        print(
+            "Max diff (HF base vs HF LoRA prefill):",
+            torch.max(torch.abs(hf_no_lora_prefill - hf_prefill)),
+        )
+
+        if hf_prefill.shape[0] <= 100:
+            assert torch.all(torch.abs(hf_prefill - srt_prefill) < prefill_tol), (
+                f"Prefill logprobs mismatch for base '{base_path}', adaptor '{adaptor.name}', "
+                f"backend '{backend}', prompt: '{prompt[:50]}...'"
+            )
+
+        if hf_decode.shape[0] <= 100:
+            assert torch.all(torch.abs(hf_decode - srt_decode) < decode_tol), (
+                f"Decode logprobs mismatch for base '{base_path}', adaptor '{adaptor.name}', "
+                f"backend '{backend}', prompt: '{prompt[:50]}...'"
+            )
+
+        if rouge_score < rouge_tol:
+
+            raise AssertionError(
+                f"ROUGE-L score {rouge_score} below tolerance {rouge_tol} "
+                f"for base '{base_path}', adaptor '{adaptor.name}', backend '{backend}', prompt: '{prompt[:50]}...'"
+            )
+
+    def run_backend_batch(
+        self,
+        prompts: List[str],
+        model_case: LoRAModelCase,
+        torch_dtype: torch.dtype,
+        max_new_tokens: int,
+        backend: str,
+    ):
+        # TODO: Implement batch processing version of run_backend
+        raise NotImplementedError(
+            "Batch processing version of run_backend is not implemented yet."
+        )
+
+    def _run_backend_on_model_cases(self, model_cases: List[LoRAModelCase]):
+        for model_case in model_cases:
+            # If skip_long_prompt is True, filter out prompts longer than 1000 characters
+            prompts = (
+                PROMPTS
+                if not model_case.skip_long_prompt
+                else [p for p in PROMPTS if len(p) < 1000]
+            )
+            for torch_dtype in TORCH_DTYPES:
+                for backend in BACKENDS:
path: test/srt/models/lora/test_lora_backend_tensor_parallel.py
user: login: Fridge003
body: Please only use `triton` backend and a fixed prompt here, otherwise there are too many tests to do. Also, here we want to test different sizes for tp, so create a `TP_SIZE` list and iterate it.
created_at: 2025-03-11T17:36:55Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,258 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import multiprocessing as mp
+import os
+import unittest
+from typing import List
+
+import torch
+from utils import BACKENDS, TORCH_DTYPES, LoRAAdaptor, LoRAModelCase
+
+from sglang.test.runners import HFRunner, SRTRunner
+from sglang.test.test_utils import calculate_rouge_l, is_in_ci
+
+CI_LORA_MODELS = [
+    LoRAModelCase(
+        base="meta-llama/Llama-3.1-8B-Instruct",
+        adaptors=[
+            LoRAAdaptor(
+                name="algoprog/fact-generation-llama-3.1-8b-instruct-lora",
+            ),
+        ],
+        max_loras_per_batch=1,
+        tp_size=2,
+    ),
+    LoRAModelCase(
+        base="meta-llama/Llama-3.1-8B-Instruct",
+        adaptors=[
+            LoRAAdaptor(
+                name="Nutanix/Meta-Llama-3.1-8B-Instruct_lora_4_alpha_16",
+                prefill_tolerance=1e-1,
+            ),
+        ],
+        max_loras_per_batch=1,
+        tp_size=2,
+    ),
+]
+
+ALL_OTHER_LORA_MODELS = [
+    LoRAModelCase(
+        base="meta-llama/Llama-2-7b-hf",
+        adaptors=[LoRAAdaptor(name="winddude/wizardLM-LlaMA-LoRA-7B")],
+        max_loras_per_batch=2,
+        tp_size=2,
+    ),
+]
+
+PROMPTS = [
+    "AI is a field of computer science focused on",
+    """
+    ### Instruction:
+    Tell me about llamas and alpacas
+    ### Response:
+    Llamas are large, long-necked animals with a woolly coat. They have two toes on each foot instead of three like other camelids (camels, dromedaries). Llamas live in the Andean mountains of South America where they graze on grasses and shrubs. Alpaca is another name for domesticated llama. The word "alpaca" comes from an Incan language meaning "golden fleece." Alpacas look very similar to llamas but are smaller than their wild relatives. Both species were used by ancient people as pack animals and for meat. Today both llamas and alpacas are raised primarily for their fiber which can be spun into yarn or knitted into clothing.
+    ### Question 2:
+    What do you know about llamas?
+    ### Answer:
+    """,
+]
+
+
+class TestLoRABackend(unittest.TestCase):
path: test/srt/models/lora/test_lora_backend_tensor_parallel.py
user: login: Fridge003
body: Change the class name to `TestLoRATP` might be better.
created_at: 2025-03-11T17:37:17Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,258 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import multiprocessing as mp
+import os
+import unittest
+from typing import List
+
+import torch
+from utils import BACKENDS, TORCH_DTYPES, LoRAAdaptor, LoRAModelCase
+
+from sglang.test.runners import HFRunner, SRTRunner
+from sglang.test.test_utils import calculate_rouge_l, is_in_ci
+
+CI_LORA_MODELS = [
+    LoRAModelCase(
+        base="meta-llama/Llama-3.1-8B-Instruct",
+        adaptors=[
+            LoRAAdaptor(
+                name="algoprog/fact-generation-llama-3.1-8b-instruct-lora",
+            ),
+        ],
+        max_loras_per_batch=1,
+        tp_size=2,
+    ),
+    LoRAModelCase(
+        base="meta-llama/Llama-3.1-8B-Instruct",
+        adaptors=[
+            LoRAAdaptor(
+                name="Nutanix/Meta-Llama-3.1-8B-Instruct_lora_4_alpha_16",
+                prefill_tolerance=1e-1,
+            ),
+        ],
+        max_loras_per_batch=1,
+        tp_size=2,
+    ),
+]
+
+ALL_OTHER_LORA_MODELS = [
+    LoRAModelCase(
+        base="meta-llama/Llama-2-7b-hf",
+        adaptors=[LoRAAdaptor(name="winddude/wizardLM-LlaMA-LoRA-7B")],
+        max_loras_per_batch=2,
+        tp_size=2,
+    ),
+]
+
+PROMPTS = [
+    "AI is a field of computer science focused on",
+    """
+    ### Instruction:
+    Tell me about llamas and alpacas
+    ### Response:
+    Llamas are large, long-necked animals with a woolly coat. They have two toes on each foot instead of three like other camelids (camels, dromedaries). Llamas live in the Andean mountains of South America where they graze on grasses and shrubs. Alpaca is another name for domesticated llama. The word "alpaca" comes from an Incan language meaning "golden fleece." Alpacas look very similar to llamas but are smaller than their wild relatives. Both species were used by ancient people as pack animals and for meat. Today both llamas and alpacas are raised primarily for their fiber which can be spun into yarn or knitted into clothing.
+    ### Question 2:
+    What do you know about llamas?
+    ### Answer:
+    """,
+]
+
+
+class TestLoRABackend(unittest.TestCase):
+    def run_backend(
path: test/srt/models/lora/test_lora_backend_tensor_parallel.py
user: login: Fridge003
body: rename to `run_tp`
created_at: 2025-03-11T17:37:41Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,258 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import multiprocessing as mp
+import os
+import unittest
+from typing import List
+
+import torch
+from utils import BACKENDS, TORCH_DTYPES, LoRAAdaptor, LoRAModelCase
+
+from sglang.test.runners import HFRunner, SRTRunner
+from sglang.test.test_utils import calculate_rouge_l, is_in_ci
+
+CI_LORA_MODELS = [
+    LoRAModelCase(
+        base="meta-llama/Llama-3.1-8B-Instruct",
+        adaptors=[
+            LoRAAdaptor(
+                name="algoprog/fact-generation-llama-3.1-8b-instruct-lora",
+            ),
+        ],
+        max_loras_per_batch=1,
+        tp_size=2,
+    ),
+    LoRAModelCase(
+        base="meta-llama/Llama-3.1-8B-Instruct",
+        adaptors=[
+            LoRAAdaptor(
+                name="Nutanix/Meta-Llama-3.1-8B-Instruct_lora_4_alpha_16",
+                prefill_tolerance=1e-1,
+            ),
+        ],
+        max_loras_per_batch=1,
+        tp_size=2,
+    ),
+]
+
+ALL_OTHER_LORA_MODELS = [
+    LoRAModelCase(
+        base="meta-llama/Llama-2-7b-hf",
+        adaptors=[LoRAAdaptor(name="winddude/wizardLM-LlaMA-LoRA-7B")],
+        max_loras_per_batch=2,
+        tp_size=2,
+    ),
+]
+
+PROMPTS = [
+    "AI is a field of computer science focused on",
+    """
+    ### Instruction:
+    Tell me about llamas and alpacas
+    ### Response:
+    Llamas are large, long-necked animals with a woolly coat. They have two toes on each foot instead of three like other camelids (camels, dromedaries). Llamas live in the Andean mountains of South America where they graze on grasses and shrubs. Alpaca is another name for domesticated llama. The word "alpaca" comes from an Incan language meaning "golden fleece." Alpacas look very similar to llamas but are smaller than their wild relatives. Both species were used by ancient people as pack animals and for meat. Today both llamas and alpacas are raised primarily for their fiber which can be spun into yarn or knitted into clothing.
+    ### Question 2:
+    What do you know about llamas?
+    ### Answer:
+    """,
+]
+
+
+class TestLoRABackend(unittest.TestCase):
+    def run_backend(
+        self,
+        prompt: str,
+        model_case: LoRAModelCase,
+        torch_dtype: torch.dtype,
+        max_new_tokens: int,
+        backend: str,
+    ):
+        """
+        Run backend tests for a single prompt and model case.
+        """
+        base_path = model_case.base
+        adaptor = model_case.adaptors[0]
+        print(
+            f"\n========== Testing backend '{backend}' for base '{base_path}' --- "
+            f"Prompt '{prompt[:50]}...' using adaptor '{adaptor.name}' ---"
+        )
+        with SRTRunner(
+            base_path,
+            torch_dtype=torch_dtype,
+            model_type="generation",
+            tp_size=model_case.tp_size,
+            lora_paths=[adaptor.name for adaptor in model_case.adaptors],
+            max_loras_per_batch=model_case.max_loras_per_batch,
+            lora_backend=backend,
+            disable_cuda_graph=True,
+            disable_radix_cache=True,
+            mem_fraction_static=0.88,
+            disable_custom_all_reduce=True,
+        ) as srt_runner:
+            srt_outputs = srt_runner.forward(
+                [prompt], max_new_tokens=max_new_tokens, lora_paths=[adaptor.name]
+            )
+
+        with HFRunner(
+            base_path, torch_dtype=torch_dtype, model_type="generation"
+        ) as hf_runner:
+            hf_outputs = hf_runner.forward(
+                [prompt], max_new_tokens=max_new_tokens, lora_paths=[adaptor.name]
+            )
+
+        with SRTRunner(
+            base_path,
+            torch_dtype=torch_dtype,
+            model_type="generation",
+            tp_size=model_case.tp_size,
+            mem_fraction_static=0.88,
+            disable_custom_all_reduce=True,
+        ) as srt_runner:
+            srt_no_lora_outputs = srt_runner.forward(
+                [prompt], max_new_tokens=max_new_tokens
+            )
+
+        with HFRunner(
+            base_path,
+            torch_dtype=torch_dtype,
+            model_type="generation",
+        ) as hf_runner:
+            hf_no_lora_outputs = hf_runner.forward(
+                [prompt], max_new_tokens=max_new_tokens
+            )
+
+        # Use individual adapter tolerances if set, otherwise use model defaults
+        prefill_tol = (
+            adaptor.prefill_tolerance
+            if adaptor.prefill_tolerance is not None
+            else model_case.prefill_tolerance
+        )
+        decode_tol = (
+            adaptor.decode_tolerance
+            if adaptor.decode_tolerance is not None
+            else model_case.decode_tolerance
+        )
+        rouge_tol = (
+            adaptor.rouge_l_tolerance
+            if adaptor.rouge_l_tolerance is not None
+            else model_case.rouge_l_tolerance
+        )
+
+        # Compare prefill stage logprobs (HF vs SRTRunner with LoRA)
+        hf_prefill = torch.tensor(hf_outputs.top_input_logprobs[0])
+        srt_prefill = torch.tensor(srt_outputs.top_input_logprobs[0])
+        max_prefill_diff = torch.max(torch.abs(hf_prefill - srt_prefill))
+        print("Max prefill diff (HF vs SRT):", max_prefill_diff)
+
+        # Compare decode stage logprobs
+        hf_decode = torch.tensor(hf_outputs.top_output_logprobs[0])
+        srt_decode = torch.tensor(srt_outputs.top_output_logprobs[0])
+        max_decode_diff = torch.max(torch.abs(hf_decode - srt_decode))
+        print("Max decode diff (HF vs SRT):", max_decode_diff)
+
+        srt_output_str = srt_outputs.output_strs[0].strip()
+        hf_output_str = hf_outputs.output_strs[0].strip()
+        rouge_score = calculate_rouge_l([srt_output_str], [hf_output_str])[0]
+        print("ROUGE-L score:", rouge_score)
+        print("SRT output:", srt_output_str)
+        print("HF output:", hf_output_str)
+
+        # Additional: compare prefill outputs between base model (no LoRA) and LoRA model for reference
+        hf_no_lora_prefill = torch.tensor(hf_no_lora_outputs.top_input_logprobs[0])
+        srt_no_lora_prefill = torch.tensor(srt_no_lora_outputs.top_input_logprobs[0])
+        print(
+            "Max diff (SRT base vs SRT LoRA prefill):",
+            torch.max(torch.abs(srt_no_lora_prefill - srt_prefill)),
+        )
+        print(
+            "Max diff (HF base vs HF LoRA prefill):",
+            torch.max(torch.abs(hf_no_lora_prefill - hf_prefill)),
+        )
+
+        if hf_prefill.shape[0] <= 100:
+            assert torch.all(torch.abs(hf_prefill - srt_prefill) < prefill_tol), (
+                f"Prefill logprobs mismatch for base '{base_path}', adaptor '{adaptor.name}', "
+                f"backend '{backend}', prompt: '{prompt[:50]}...'"
+            )
+
+        if hf_decode.shape[0] <= 100:
+            assert torch.all(torch.abs(hf_decode - srt_decode) < decode_tol), (
+                f"Decode logprobs mismatch for base '{base_path}', adaptor '{adaptor.name}', "
+                f"backend '{backend}', prompt: '{prompt[:50]}...'"
+            )
+
+        if rouge_score < rouge_tol:
+
+            raise AssertionError(
+                f"ROUGE-L score {rouge_score} below tolerance {rouge_tol} "
+                f"for base '{base_path}', adaptor '{adaptor.name}', backend '{backend}', prompt: '{prompt[:50]}...'"
+            )
+
+    def run_backend_batch(
path: test/srt/models/lora/test_lora_backend_tensor_parallel.py
user: login: Fridge003
body: rename to `run_tp_batch`
created_at: 2025-03-11T17:37:58Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,258 @@
+# Copyright 2023-2024 SGLang Team
path: test/srt/models/lora/test_lora_backend_tensor_parallel.py
user: login: Fridge003
body: rename the file to `test_lora_tp.py` would be better
created_at: 2025-03-11T17:41:56Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,258 @@
+# Copyright 2023-2024 SGLang Team
path: test/srt/models/lora/test_lora_backend_tensor_parallel.py
user: login: aoshen524
body: sure.
created_at: 2025-03-11T18:30:47Z
author_association: CONTRIBUTOR

PR comments:
user: login: Fridge003
created_at: 2025-03-11T01:06:59Z
body: Great work! Also, to ensure the performance of lora after tp, please paste some benchmark results before/after enabling tp. You can refer to the benchmark in #3161 for example.

author_association: COLLABORATOR
user: login: aoshen524
created_at: 2025-03-11T01:41:25Z
body: > Great work! Also, to ensure the performance of lora after tp, please paste some benchmark results before/after enabling tp. You can refer to the benchmark in #3161 for example.

Good advice. But no extra communication kernel launch or used when supporting lora tp. Still recommend to do it?
author_association: CONTRIBUTOR
user: login: Fridge003
created_at: 2025-03-11T02:43:08Z
body: > > Great work! Also, to ensure the performance of lora after tp, please paste some benchmark results before/after enabling tp. You can refer to the benchmark in #3161 for example.
> 
> Good advice. But no extra communication kernel launch or used when supporting lora tp. Still recommend to do it?

Yes, just make sure lora with tp is not too slow.
author_association: COLLABORATOR
user: login: aoshen524
created_at: 2025-03-11T14:20:35Z
body: > > > Great work! Also, to ensure the performance of lora after tp, please paste some benchmark results before/after enabling tp. You can refer to the benchmark in #3161 for example.
> > 
> > 
> > Good advice. But no extra communication kernel launch or used when supporting lora tp. Still recommend to do it?
> 
> Yes, just make sure lora with tp is not too slow.

sure
author_association: CONTRIBUTOR