Use model fast forward

Author: jlamypoirier

scripts/run_all_benchmark_breakdown.sh

@@ -9,12 +9,12 @@
 ./scripts/run_benchmark_breakdown.sh santacoder bigcode/gpt_bigcode-santacoder 256 2040 11 1 v2_
 
 # Large model
-./scripts/run_benchmark_breakdown.sh large_model ./data/bigcode_large-model 1 8190 11 0 v2_
-./scripts/run_benchmark_breakdown.sh large_model ./data/bigcode_large-model 8 8190 11 0 v2_
-./scripts/run_benchmark_breakdown.sh large_model ./data/bigcode_large-model 32 8190 11 0 v2_
-./scripts/run_benchmark_breakdown.sh large_model ./data/bigcode_large-model 256 8190 11 0  v2_ # OOM?
+./scripts/run_benchmark_breakdown.sh starcoder ./data/bigcode_large-model 1 8190 11 0 v2_
+./scripts/run_benchmark_breakdown.sh starcoder ./data/bigcode_large-model 8 8190 11 0 v2_
+./scripts/run_benchmark_breakdown.sh starcoder ./data/bigcode_large-model 32 8190 11 0 v2_
+./scripts/run_benchmark_breakdown.sh starcoder ./data/bigcode_large-model 256 8190 11 0  v2_ # OOM?
 
-./scripts/run_benchmark_breakdown.sh large_model ./data/bigcode_large-model 1 8190 29 1 v2_ 1
-./scripts/run_benchmark_breakdown.sh large_model ./data/bigcode_large-model 8 8190 29 1 v2_ 1
-./scripts/run_benchmark_breakdown.sh large_model ./data/bigcode_large-model 32 8190 29 1 v2_ 1
-./scripts/run_benchmark_breakdown.sh large_model ./data/bigcode_large-model 256 8190 29 1  v2_ 1 # OOM?
+./scripts/run_benchmark_breakdown.sh starcoder ./data/bigcode_large-model 1 8190 29 1 v2_ 1
+./scripts/run_benchmark_breakdown.sh starcoder ./data/bigcode_large-model 8 8190 29 1 v2_ 1
+./scripts/run_benchmark_breakdown.sh starcoder ./data/bigcode_large-model 32 8190 29 1 v2_ 1
+./scripts/run_benchmark_breakdown.sh starcoder ./data/bigcode_large-model 256 8190 29 1  v2_ 1 # OOM?

scripts/run_textgen_benchmark_breakdown.sh

@@ -14,13 +14,13 @@ CYCLES=${7:-10}
 
 SAVE_DIR=data/benchmarks/v3
 #BATCH_SIZES="1 2 4 8 16 24 32 48 64 96 128 160 224 256"
-RUN="python3 src/main.py --pipeline_class=TG_Pipeline --max_log_outputs=0 --dtype=float16 --device=cuda  --custom_generate  --breakdown_latency --ignore_oom --no_fast_init "
+RUN="python3 -m src.main --pipeline_class=TG_Pipeline --max_log_outputs=0 --dtype=float16 --device=cuda  --custom_generate  --breakdown_latency --ignore_oom --no_fast_init "
 
 
-IMPL=("flash" "santa" "causal" "vector" "bigcode")
+IMPL=("flash" "causal" "vector" "bigcode")
 
 
-STEP=("--no_prefill" "--no_cache")
+STEP=("" "--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"
@@ -31,13 +31,14 @@ run () { # run(step, runtime, attn)
   then
     echo "Skipping existing $FILE_NAME"
   else
-    CMD="MODEL_TYPE=${IMPL[$2]} $RUN $COMMON ${STEP[$1]} --save=$FILE_NAME"
-    echo "$CMD"
+    export MODEL_TYPE="${IMPL[$2]}"
+    CMD="$RUN $COMMON ${STEP[$1]} --save=$FILE_NAME"
+    echo "MODEL_TYPE=${IMPL[$2]} $CMD"
     $CMD
   fi
 }
 
-for impl in {0..4}
+for impl in {0..3}
 do
   if [ "${STEP_ID}" -eq "0" ]
   then

src/parse_breakdown_results.py

@@ -62,8 +62,8 @@ def main(argv: Optional[List[str]] = None) -> None:
     dirname = args.input_dir.stem
     if title is None:
         try:
-            name, _, bs, _, _, _, _, step = dirname.rsplit("_", 7)
-            title = f"{name} {step}, bs = {bs}"
+            name, _, bs, _, _, _, _, step, cycles = dirname.rsplit("_", 8)
+            title = f"{name}, bs = {bs} (s={step}, c={cycles})"
         except ValueError:
             title = dirname
 

src/pipeline.py

@@ -619,8 +619,10 @@ def _generate_textgen(
         with torch.inference_mode():
             for key_length in range(input_length, output_length, key_length_step):
                 try:
-                    if key_length_step > 1 or not use_cache or not do_prefill:
-                        self._update_generate_batch(batch, use_cache, do_prefill, key_length)
+                    if (key_length_step > 1 and key_length>key_length) or not use_cache or not do_prefill:
+                        if not hasattr(self.model,"fast_forward"):
+                            raise NotImplementedError()
+                        self.model.fast_forward(batch, key_length, use_cache)
                         last_time = self._get_time(breakdown_latency)
                     generated, batch = self.model.generate_token(batch)
                     t2 = self._get_time(breakdown_latency)