Fixed issues

server/lorax_server/adapters/weights.py

@@ -118,7 +118,7 @@ class AdapterBatchData:
     data: Dict[str, Dict[str, BatchAdapterWeights]]
 
     # layer type -> fused lora weights
-    layer_to_lora_weights: Dict[str, Tuple[torch.Tensor, torch.Tensor]]
+    layer_to_lora_weights: Dict[Tuple[str, int], Tuple[torch.Tensor, torch.Tensor]]
 
     punica_wrapper: "PunicaWrapper"
 
@@ -128,6 +128,7 @@ class AdapterBatchData:
     def from_meta(
         meta: AdapterBatchMetadata,
         weights: Dict[str, LayerAdapterWeights],
+        layer_to_lora_weights: Dict[Tuple[str, int], Tuple[torch.Tensor, torch.Tensor]],
         punica_wrapper: "PunicaWrapper",
         prefill: bool,
         prefill_head_indices: Optional[torch.Tensor],
@@ -139,7 +140,13 @@ def from_meta(
             layer_weights = v.get_data(meta, k, prefill, prefill_head_indices if k == LM_HEAD else None)
             if layer_weights:
                 data[k] = layer_weights
-        return AdapterBatchData(meta=meta, data=data, punica_wrapper=punica_wrapper, prefill=prefill)
+        return AdapterBatchData(
+            meta=meta, 
+            data=data, 
+            layer_to_lora_weights=layer_to_lora_weights, 
+            punica_wrapper=punica_wrapper, 
+            prefill=prefill,
+        )
 
     def ranks(self) -> Set[int]:
         # TODO(travis): refactor to be less coupled to lora implementation

server/lorax_server/models/custom_modeling/flash_qwen2_modeling.py

@@ -303,8 +303,8 @@ def __init__(self, prefix, config, weights, layer_id):
             layer_id,
             [MLP_GATE_PROJ, MLP_UP_PROJ],
             sizes=[
-                config.intermediate_size // 2,
-                config.intermediate_size // 2,
+                config.intermediate_size,
+                config.intermediate_size,
             ],
             process_group=weights.process_group,
         )

server/lorax_server/models/model.py

@@ -5,6 +5,7 @@
 
 from lorax_server.adapters.lora import LoraWeights
 from lorax_server.adapters.medusa_lora import MedusaLoraWeights
+from lorax_server.utils.sgmv import pad_to_min_rank
 import torch
 from loguru import logger
 from transformers import PreTrainedTokenizerBase
@@ -252,14 +253,14 @@ def register_preloaded_adapters(
         self.preloaded_adapters.extend(preloaded_adapters)
 
         # For Triton kernels: need weights into contiguous tensor
-        # dict of layer_name -> (lora_a_weights, lora_b_weights)
+        # dict of (layer_name, layer_id) -> (lora_a_weights, lora_b_weights)
         # where:
         #   lora_a_weights = [num_adapters, r, hidden_size] 
         #   lora_b_weights = [num_adapters, hidden_size, r]
         self.layer_to_lora_weights = {}
         for layer_name, layer_adapter_weights in self.layer_to_adapter_weights.items():
-            lora_a_weights = []
-            lora_b_weights = []
+            layer_id_to_lora_a_weights = defaultdict(list)
+            layer_id_to_lora_b_weights = defaultdict(list)
             for i, adapter in enumerate(preloaded_adapters):
                 adapter_index = adapter.adapter_index
                 adapter_weights = layer_adapter_weights.adapter_weights.get(adapter_index)
@@ -271,17 +272,31 @@ def register_preloaded_adapters(
                         # only applicable to lora for now
                         continue
 
-                # transpose to ensure col major
-                lora_a = adapter_weights.weights_a_t
-                lora_b = adapter_weights.weights_b_t
-
-                lora_a_weights.append(lora_a)
-                lora_b_weights.append(lora_b)
+                # transpose into col major
+                lora_a = adapter_weights.weights_a.transpose(1, 2)
+                lora_b = adapter_weights.weights_b.transpose(1, 2)
+
+                nlayers = lora_a.size(0)
+                for layer_id in range(nlayers):
+                    layer_id_to_lora_a_weights[layer_id].append(lora_a[layer_id])
+                    layer_id_to_lora_b_weights[layer_id].append(lora_b[layer_id])
 
-            # stack into [num_adapters, r, hidden_size] and [num_adapters, hidden_size, r]
-            lora_a_weights = torch.stack(lora_a_weights, device=self.device).contiguous()
-            lora_b_weights = torch.stack(lora_b_weights, device=self.device).contiguous()
-            self.layer_to_lora_weights[layer_name] = (lora_a_weights, lora_b_weights)
+            for layer_id, lora_a_weights in layer_id_to_lora_a_weights.items():
+                lora_b_weights = layer_id_to_lora_b_weights[layer_id]
+
+                # right pad every adapter to the max rank
+                # TODO(travis)
+                # r = max([w.size(-1) for w in lora_b_weights])
+                # lora_a_weights = [pad_to_min_rank(w, 1, r) for w in lora_a_weights]
+                # lora_b_weights = [pad_to_min_rank(w, 2, r) for w in lora_b_weights]
+
+                # stack into [num_adapters, r, hidden_size] and [num_adapters, hidden_size, r]
+                lora_a_weights = torch.stack(lora_a_weights).to(self.device).contiguous()
+                lora_b_weights = torch.stack(lora_b_weights).to(self.device).contiguous()
+                print("!!! lora_a_weights", lora_a_weights.shape, layer_name, layer_id)
+                print("!!! lora_b_weights", lora_b_weights.shape)
+                # ('self_attn.q_proj', 32)
+                self.layer_to_lora_weights[(layer_name, layer_id)] = (lora_a_weights, lora_b_weights)
 
     def load_adapter(
         self,

server/lorax_server/utils/layers.py

@@ -76,6 +76,7 @@ def forward_layer_type(
         data: Optional["BatchLoraWeights"] = data.get(LORA) if data is not None else None
 
         if has_sgmv() and data is not None and data.can_vectorize(self.process_group):
+            print("!!! layer_type", layer_type, "start_idx", start_idx, "end_idx", end_idx, "result", result.shape)
             if end_idx - start_idx != result.shape[1]:
                 # proj = torch.zeros_like(result[:, start_idx:end_idx])
                 y_offset = start_idx
@@ -89,7 +90,7 @@ def forward_layer_type(
                 # lora_a_ptr = rank_segments.lora_a_ptr
                 # lora_b_ptr = rank_segments.lora_b_ptr
 
-                lora_a_weights, lora_b_weights = adapter_data.layer_to_lora_weights[layer_type]
+                lora_a_weights, lora_b_weights = adapter_data.layer_to_lora_weights[(layer_type, self.layer_id)]
                 adapter_data.punica_wrapper.add_lora(
                     result,
                     input,
@@ -230,6 +231,7 @@ def forward(self, input: torch.Tensor, adapter_data: "AdapterBatchData") -> torc
                 end_idx = offset // self.process_group.size()
             else:
                 end_idx = result.shape[1]
+            print("!!! sizes", self.sizes, self.process_group.size())
 
             result = self.forward_layer_type(result, input, adapter_data, layer_name, start_idx, end_idx)
 

server/lorax_server/utils/ops/sgmv_expand.py

@@ -9,7 +9,7 @@
 import triton
 import triton.language as tl
 
-from lorax_server.utils.ops import libentry
+from lorax_server.utils.ops.libentry import libentry
 
 
 @libentry()
@@ -128,6 +128,14 @@ def sgmv_expand(
         add_inputs (bool, optional):  Defaults to False. adds the final lora 
             results to the output.
     """
+    print("!!! inputs", inputs.shape)
+    print("!!! lora_b_weights", lora_b_weights.shape)
+    print("!!! output_tensor", output_tensor.shape)
+    print("!!! b_seq_start_loc", b_seq_start_loc)
+    print("!!! seq_len_tensor", seq_len_tensor)
+    print("!!! lora_indices_tensor", lora_indices_tensor)
+    print("!!! batches", batches)
+    print("!!! max_seq_length", max_seq_length)
 
     assert inputs.dtype in [torch.float16, torch.bfloat16, torch.float32]
     assert lora_b_weights.dtype in [

server/lorax_server/utils/ops/sgmv_expand_slice.py

@@ -9,7 +9,7 @@
 import triton
 import triton.language as tl
 
-from lorax_server.utils.ops import libentry
+from lorax_server.utils.ops.libentry import libentry
 
 
 @libentry()
@@ -137,6 +137,16 @@ def sgmv_expand_slice(
         add_inputs (bool, optional):  Defaults to False. adds the final lora 
             results to the output..
     """
+    print("!!! inputs", inputs.shape)
+    print("!!! lora_b_weights", lora_b_weights.shape)
+    print("!!! output_tensor", output_tensor.shape)
+    print("!!! b_seq_start_loc", b_seq_start_loc)
+    print("!!! seq_len_tensor", seq_len_tensor)
+    print("!!! lora_indices_tensor", lora_indices_tensor)
+    print("!!! batches", batches)
+    print("!!! max_seq_length", max_seq_length)
+    print("!!! slice_offset", slice_offset)
+    print("!!! slice_size", slice_size)
 
     assert inputs.dtype in [torch.float16, torch.bfloat16, torch.float32]
     assert lora_b_weights.dtype in [

server/lorax_server/utils/ops/sgmv_shrink.py

@@ -9,7 +9,7 @@
 import triton
 import triton.language as tl
 
-from lorax_server.utils.ops import libentry
+from lorax_server.utils.ops.libentry import libentry
 
 
 @libentry()
@@ -131,6 +131,15 @@ def sgmv_shrink(
             in the batch
         scaling (float):  Scaling factor.
     """
+    print("!!! inputs", inputs.shape)
+    print("!!! lora_a_weights", lora_a_weights.shape)
+    print("!!! output_tensor", output_tensor.shape)
+    print("!!! b_seq_start_loc", b_seq_start_loc)
+    print("!!! seq_len_tensor", seq_len_tensor)
+    print("!!! lora_indices_tensor", lora_indices_tensor)
+    print("!!! batch_size", batches)
+    print("!!! max_seq_length", max_seq_length)
+    print("!!! scaling", scaling)
     assert inputs.dtype == lora_a_weights.dtype
     assert inputs.dtype in [torch.float16, torch.bfloat16]
     assert lora_a_weights.dtype in [

server/lorax_server/utils/sgmv.py

@@ -3,10 +3,16 @@
 from functools import lru_cache
 from typing import TYPE_CHECKING, Callable, List, Optional, Tuple, Union
 
-from lorax_server.utils.ops import bgmv_expand, bgmv_expand_slice, bgmv_shrink, sgmv_expand, sgmv_expand_slice, sgmv_shrink
 import torch
 import torch.nn.functional as F
 
+from lorax_server.utils.ops.bgmv_expand import bgmv_expand
+from lorax_server.utils.ops.bgmv_expand_slice import bgmv_expand_slice
+from lorax_server.utils.ops.bgmv_shrink import bgmv_shrink
+from lorax_server.utils.ops.sgmv_expand import sgmv_expand
+from lorax_server.utils.ops.sgmv_expand_slice import sgmv_expand_slice
+from lorax_server.utils.ops.sgmv_shrink import sgmv_shrink
+
 if TYPE_CHECKING:
     from lorax_server.adapters.weights import AdapterBatchMetadata
 
@@ -39,7 +45,10 @@ def pad_rank(t: torch.Tensor, dim: int, world_size: int) -> torch.Tensor:
     # tensor parallelism will result in effective rank being divided by world_size,
     # so we need to scale the min rank to offset that effect
     min_rank = MIN_SGMV_RANK * world_size
+    return pad_to_min_rank(t, dim, min_rank)
+
 
+def pad_to_min_rank(t: torch.Tensor, dim: int, min_rank: int) -> torch.Tensor:
     # if we're at or below the min rank, pad up to the min rank
     # otherwise, pad to the nearest multiple of the block size
     current_rank = t.size(dim)