bugfix

Author: james

README.md

@@ -94,7 +94,7 @@ CUDA_VISIBLE_DEVICES=0 python generate.py \
     --interactive
 ```
 
-4. Run sparse inference (`scripts/run.sh`)!:
+4. Run sparse inference! (`scripts/run.sh`):
 ```bash
 CUDA_VISIBLE_DEVICES=0 python generate.py \
     --compile \ 

gpt-fast/scripts/speculate_34B_bf16.sh

@@ -0,0 +1 @@
+time torchrun --standalone --nproc_per_node=4 generate.py --compile --checkpoint_path $OUTPUT_PATH/meta-llama/Llama-2-7b-hf/model.pth --hist_path ../models/Llama-2-70B/histograms --sparsity 0.5 --prompt "Hello, my name is "

gpt-fast/scripts/speculate_70B_int4.sh

@@ -14,24 +14,32 @@ def init_func(nargs):
 # NOTE: will need to warm up kernels each time, triton autotune caching isn't a thing right now
 
 configs=[
+    triton.Config({"BLOCK_M": 64, "BLOCK_N": 128}, num_warps=2, pre_hook=init_to_zero("Y")), 
+
+    triton.Config({"BLOCK_M": 64, "BLOCK_N": 64}, num_warps=4, pre_hook=init_to_zero("Y")),
     triton.Config({"BLOCK_M": 8, "BLOCK_N": 128}, num_warps=2, pre_hook=init_to_zero("Y")),
     triton.Config({"BLOCK_M": 16, "BLOCK_N": 256}, num_warps=4, pre_hook=init_to_zero("Y")),
     triton.Config({"BLOCK_M": 16, "BLOCK_N": 256}, num_warps=4, pre_hook=init_to_zero("Y")),
-    triton.Config({"BLOCK_M": 16, "BLOCK_N": 512}, num_warps=4, pre_hook=init_to_zero("Y")),
-    #triton.Config({"BLOCK_M": 16, "BLOCK_N": 1024}, num_warps=4, pre_hook=init_to_zero("Y")),
     triton.Config({"BLOCK_M": 32, "BLOCK_N": 256}, num_warps=4, pre_hook=init_to_zero("Y")),
-    triton.Config({"BLOCK_M": 32, "BLOCK_N": 512}, num_warps=4, pre_hook=init_to_zero("Y")),
-    #triton.Config({"BLOCK_M": 32, "BLOCK_N": 1024}, num_warps=4, pre_hook=init_to_zero("Y")),
     triton.Config({"BLOCK_M": 64, "BLOCK_N": 256}, num_warps=4, pre_hook=init_to_zero("Y")),
-    triton.Config({"BLOCK_M": 64, "BLOCK_N": 512}, num_warps=4, pre_hook=init_to_zero("Y")),
-    #triton.Config({"BLOCK_M": 64, "BLOCK_N": 1024}, num_warps=4, pre_hook=init_to_zero("Y")),
     triton.Config({"BLOCK_M": 128, "BLOCK_N": 16}, num_warps=4, pre_hook=init_to_zero("Y")),
     triton.Config({"BLOCK_M": 128, "BLOCK_N": 32}, num_warps=4, pre_hook=init_to_zero("Y")),
     triton.Config({"BLOCK_M": 128, "BLOCK_N": 64}, num_warps=4, pre_hook=init_to_zero("Y")),
     triton.Config({"BLOCK_M": 128, "BLOCK_N": 128}, num_warps=4, pre_hook=init_to_zero("Y")),
     triton.Config({"BLOCK_M": 128, "BLOCK_N": 256}, num_warps=4, pre_hook=init_to_zero("Y")),
-    # triton.Config({"BLOCK_M": 128, "BLOCK_N": 512}, num_warps=4, pre_hook=init_to_zero("Y")),
-    #triton.Config({"BLOCK_M": 128, "BLOCK_N": 1024}, num_warps=4, pre_hook=init_to_zero("Y")),
+
+    triton.Config({"BLOCK_M": 128, "BLOCK_N": 512}, num_warps=4, pre_hook=init_to_zero("Y")),
+    triton.Config({"BLOCK_M": 64, "BLOCK_N": 512}, num_warps=4, pre_hook=init_to_zero("Y")),
+    triton.Config({"BLOCK_M": 32, "BLOCK_N": 512}, num_warps=4, pre_hook=init_to_zero("Y")),
+    triton.Config({"BLOCK_M": 16, "BLOCK_N": 512}, num_warps=4, pre_hook=init_to_zero("Y")),
+
+
+    # Llama 3 variants can use BLOCK_N >= 1024
+    # triton.Config({"BLOCK_M": 128, "BLOCK_N": 1024}, num_warps=4, pre_hook=init_to_zero("Y")),
+    # triton.Config({"BLOCK_M": 16, "BLOCK_N": 1024}, num_warps=4, pre_hook=init_to_zero("Y")),
+    # triton.Config({"BLOCK_M": 64, "BLOCK_N": 1024}, num_warps=4, pre_hook=init_to_zero("Y")),
+    # triton.Config({"BLOCK_M": 32, "BLOCK_N": 1024}, num_warps=4, pre_hook=init_to_zero("Y")),
+    # triton.Config({"BLOCK_M": 16, "BLOCK_N": 1024}, num_warps=4, pre_hook=init_to_zero("Y")),
 ]
 
 @triton.autotune(
@@ -287,7 +295,6 @@ def forward(
         sparsity_bin: int,
         kv_size: int
     ) -> torch.Tensor:
-        return torch.matmul(x, weight.T)
         return qkv_gemv(x, weight, threshold_q, threshold_k, threshold_v, sparsity_bin, kv_size) if x.shape[1] == 1 else torch.matmul(x, weight.T)
 
 # for testing purposes, to see if overhead at 0% is really due to strengthening torch.matmul (seems like it is)

gpt-fast/scripts/speculate_7B_int4.sh

@@ -52,13 +52,16 @@
 for sample in tqdm(dataset):
     text += sample["text"] + "\n\n"
 
-
+print(len(text))
 bsz, seq_len = 10, 2048
 
+input_ids = []
+for i in range(0, len(text), seq_len):
+    ttext = text[i:i+seq_len]
+    encodings = tokenizer(ttext, truncation=True, return_tensors="pt", max_length=seq_len, return_overflowing_tokens=True, padding="max_length")
+    input_ids.append(encodings.input_ids)
 
-encodings = tokenizer(text, truncation=True, return_tensors="pt", max_length=seq_len, return_overflowing_tokens=True, padding="max_length")
-
-input_ids = encodings.input_ids[:bsz,:].to(device="cuda:0")
+input_ids = torch.cat(input_ids, dim=0)[:bsz,:].to(device="cuda:0")
 print(input_ids.shape)
 
 hidden_states = model.model.embed_tokens(input_ids)