Determine trace set from preloaded adapter set

predibase · Nov 1, 2024 · 1c70ec6 · 1c70ec6
1 parent 43c129b
commit 1c70ec6
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 2 deletions.
diff --git a/server/lorax_server/models/flash_causal_lm.py b/server/lorax_server/models/flash_causal_lm.py
@@ -1389,7 +1389,7 @@ def warmup(self, batch: FlashCausalLMBatch, max_new_tokens: int, embedding_model
                 self.device,
                 self.kv_cache,
                 self.adapter_layers,
-                self.default_traced_adapter_layers,
+                self.traced_adapter_layers,
                 self._forward_context,
                 max_total_tokens,
                 self.num_heads,

diff --git a/server/lorax_server/models/model.py b/server/lorax_server/models/model.py
@@ -233,6 +233,12 @@ def adapter_target_to_layer(self) -> Dict[str, Tuple[str, torch.Tensor]]:
     @property
     def adapter_layers(self) -> List[str]:
         return []
+
+    @property
+    def traced_adapter_layers(self) -> List[str]:
+        if self.layer_to_adapter_weights:
+            return list(self.layer_to_adapter_weights.keys())
+        return self.default_traced_adapter_layers
 
     @property
     def default_traced_adapter_layers(self) -> List[str]:
@@ -279,7 +285,7 @@ def register_preloaded_adapters(
         for layer_name, layer_adapter_weights in self.layer_to_adapter_weights.items():
             layer_id_to_lora_a_weights = defaultdict(list)
             layer_id_to_lora_b_weights = defaultdict(list)
-            for i, adapter in enumerate(preloaded_adapters):
+            for adapter in preloaded_adapters:
                 adapter_index = adapter.adapter_index
                 adapter_weights = layer_adapter_weights.adapter_weights.get(adapter_index)
                 if not isinstance(adapter_weights, LoraWeights):

diff --git a/server/lorax_server/utils/graph.py b/server/lorax_server/utils/graph.py
@@ -561,6 +561,7 @@ def get_estimated_cache_memory(self) -> int:
     def warmup(self):
         ngraphs = len(CACHED_BATCH_SIZES) * len(CACHED_MAX_RANKS)
         pool = None
+        logger.info("Tracing CUDA graphs with initial adapter layers: {}", self.default_traced_adapter_layers)
         with tqdm(total=ngraphs, desc="Trace CUDA graphs") as pbar:
             for batch_size in reversed(CACHED_BATCH_SIZES):
                 pbar.set_postfix({"batch_size": batch_size})