stuff

Author: jlamypoirier

scripts/run_benchmark_breakdown.sh

@@ -56,7 +56,9 @@ run () { # run(step, runtime, attn)
   then
     echo "Skipping existing $FILE_NAME"
   else
-    $RUN $COMMON  ${RUNTIME[$2]} ${ATTN[$3]} ${STEP[$1]} --save="$FILE_NAME"
+    CMD="$RUN $COMMON  ${RUNTIME[$2]} ${ATTN[$3]} ${STEP[$1]} --save=$FILE_NAME"
+    echo "$CMD"
+    $CMD
   fi
 }
 

scripts/run_textgen_benchmark_breakdown.sh

@@ -0,0 +1,56 @@
+
+# Santacoder prefill.
+# ./scripts/run_benchmark_breakdown.sh santacoder bigcode/gpt_bigcode-santacoder 32 2040 5 0
+# Santacoder decode (fewer data points because slower)
+# ./scripts/run_benchmark_breakdown.sh santacoder bigcode/gpt_bigcode-santacoder 32 2040 11 1
+MODEL_NAME=${1:-"santacoder"}
+MODEL_PATH=${2:-"bigcode/gpt_bigcode-santacoder"}
+BATCH_SIZE=${3:-32}
+MAX_NEW_TOKENS=${4:-2040}
+# Prime number to see key length padding effect.
+TOKEN_STEP=${5:-5}
+STEP_ID=${6:-""}
+FILE_PREFIX=${7:-""}
+CYCLES=${8:-10}
+
+SAVE_DIR=data/benchmarks/v2
+#BATCH_SIZES="1 2 4 8 16 24 32 48 64 96 128 160 224 256"
+RUN="python3 src/main.py --max_log_outputs=0 --dtype=float16 --device=cuda  --custom_generate  --breakdown_latency --ignore_oom"
+
+
+RUNTIME=("")
+RUNTIME_NAMES=("base")
+
+ATTN=( \
+  "--pipeline_class=TG_Pipeline" \
+  )
+ATTN_NAME=( \
+  "textgen" \
+  )
+
+
+STEP=("--no_prefill" "--no_cache")
+STEP_NAME=("decode" "prefill")
+
+COMMON="--pretrained_model=$MODEL_PATH --tokenizer=$MODEL_PATH --cycles=$CYCLES --max_input_length=1 --max_new_tokens=$MAX_NEW_TOKENS --key_length_step=$TOKEN_STEP --batch_size=$BATCH_SIZE predict_last_token=True"
+
+run () { # run(step, runtime, attn)
+  FILE_NAME="$SAVE_DIR"/"$MODEL_NAME"_bs_"$BATCH_SIZE"_tok_"$MAX_NEW_TOKENS"_step_"$TOKEN_STEP"_"${STEP_NAME[$1]}"/"$FILE_PREFIX""${RUNTIME_NAMES[$2]}"_"${ATTN_NAME[$3]}".json
+  if [ -f "$FILE_NAME" ];
+  then
+    echo "Skipping existing $FILE_NAME"
+  else
+    CMD="$RUN $COMMON  ${RUNTIME[$2]} ${ATTN[$3]} ${STEP[$1]} --save=$FILE_NAME"
+    echo "$CMD"
+    $CMD
+  fi
+}
+
+if [ "${STEP_ID}" -eq "0" ]
+then
+  # Decode (default attn only)
+  run 0 0 0
+else
+  # Prefill
+  run 1 0 0
+fi

src/pipeline.py

@@ -21,6 +21,10 @@
     GPTBigCodeConfig,
 )
 
+from transformers.modeling_outputs import (
+    CausalLMOutputWithCrossAttentions,
+)
+
 
 logger = logging.getLogger(__name__)
 
@@ -413,7 +417,7 @@ def __init__(self, **kwargs):
 
         super().__init__(**kwargs)
 
-        if self.device != torch.device("cuda"):
+        if self.device != torch.device("cuda:0"):
             raise ValueError(f"Deepspeed does not support device {self.device}")
 
         if self.dtype not in (torch.float32, torch.float16, torch.bfloat16):
@@ -433,10 +437,21 @@ def __init__(self, **kwargs):
 
 class TextGenModelWrapper:
     def __init__(self, model):
+        from text_generation_server.models import CausalLM, FlashCausalLM
+
         self.model = model
+        if isinstance(self.model, FlashCausalLM):
+            self._is_flash = True
+        elif isinstance(self.model, CausalLM):
+            self._is_flash = False
+        else:
+            raise NotImplementedError()
 
     def parameters(self):
-        return self.model.parameters()
+        return []
+
+    def eval(self):
+        pass
 
     def __call__(
         self,
@@ -447,16 +462,35 @@ def __call__(
         return_dict,
         use_cache,
     ):
-        return self.model(input_ids, attention_mask, position_ids, past_key_values)
+        if self._is_flash:
+            raise NotImplementedError()
+            logits, past_key_values = self.model.forward(
+                input_ids,
+                position_ids,
+                cu_seqlens,
+                max_s,
+                past_key_values,
+                pre_allocate_past_size,
+            )
+        else:
+            logits, past_key_values = self.model.forward(input_ids, attention_mask, position_ids, past_key_values)
+        return CausalLMOutputWithCrossAttentions(
+            loss=None,
+            logits=logits,
+            past_key_values=past_key_values,
+            hidden_states=None,
+            attentions=None,
+            cross_attentions=None,
+        )
 
 
 class TG_Pipeline(Pipeline):
     def __init__(self, **kwargs):
-        if self.device != torch.device("cuda"):
-            raise ValueError(f"Textgen does not support device {self.device}")
-
         super().__init__(**kwargs)
 
+        if self.device != torch.device("cuda:0"):
+            raise ValueError(f"Textgen does not support device {self.device}")
+
     def _get_config(
         self,
         model_type: Optional[str],
@@ -475,7 +509,7 @@ def _save_pretrained(self, pretrained_model: str):
         raise NotImplementedError()
 
     def _load_pretrained(self, pretrained_model: str):
-        from text_generation_server import get_model
+        from text_generation_server.models import get_model
 
         pretrained_model, revision = parse_revision(pretrained_model)
         return TextGenModelWrapper(get_model(pretrained_model, revision, False, False))

src/utils.py

@@ -149,8 +149,11 @@ def get_inputs_from_tokens(tokens, length, tokenizer):
     raise RuntimeError("Failed to generate stable input sequences")
 
 
-def get_random_inputs(length, tokenizer, random_state):
-    return get_inputs_from_tokens(random_state.randint(0, tokenizer.vocab_size, length).tolist(), length, tokenizer)
+def get_random_inputs(lengths, tokenizer, random_state):
+    return [
+        get_inputs_from_tokens(random_state.randint(0, tokenizer.vocab_size, length).tolist(), length, tokenizer)
+        for length in lengths
+    ]
 
 
 def get_inputs_from_files(files: List[Path], lengths, tokenizer, random_state):