diff --git a/.cache/torch/comm_lib_trace_rank_0 b/.cache/torch/comm_lib_trace_rank_0 new file mode 100644 index 0000000000000000000000000000000000000000..0d806759497252c1149322598775e6325d84e359 Binary files /dev/null and b/.cache/torch/comm_lib_trace_rank_0 differ diff --git a/.cache/torch/comm_lib_trace_rank_1 b/.cache/torch/comm_lib_trace_rank_1 new file mode 100644 index 0000000000000000000000000000000000000000..0d806759497252c1149322598775e6325d84e359 Binary files /dev/null and b/.cache/torch/comm_lib_trace_rank_1 differ diff --git a/.conda/aau_token b/.conda/aau_token new file mode 100644 index 0000000000000000000000000000000000000000..3efc135b3ac2062c15b92a99ecabdd7dc545313a --- /dev/null +++ b/.conda/aau_token @@ -0,0 +1 @@ +YN_RyWTyaweE0R_BuNYxb- \ No newline at end of file diff --git a/.conda/aau_token_host b/.conda/aau_token_host new file mode 100644 index 0000000000000000000000000000000000000000..ab11a34065f0107321da54ad5f3240587563a182 --- /dev/null +++ b/.conda/aau_token_host @@ -0,0 +1 @@ +zHxE_XAQ \ No newline at end of file diff --git a/LSAQ_CoreCode/lsaq_quant.py b/LSAQ_CoreCode/lsaq_quant.py new file mode 100644 index 0000000000000000000000000000000000000000..494aebc03cda552c9d68d86d1ec6e842a054f091 --- /dev/null +++ b/LSAQ_CoreCode/lsaq_quant.py @@ -0,0 +1,170 @@ +import os +import torch +import torch.nn as nn +import numpy as np +from transformers import AutoTokenizer, AutoModelForCausalLM +import tqdm +import json +import math +import torch.nn.functional as F + + +from datasets import load_dataset + +@torch.no_grad() +def quantize_weight_per_channel_absmax(w, n_bits=8): + # w: (out_features, in_features) + scales = w.abs().max(dim=-1, keepdim=True)[0] + q_max = 2 ** (n_bits - 1) - 1 + scales.clamp_(min=1e-5).div_(q_max) + w.div_(scales).round_().mul_(scales) + return w + + +@torch.no_grad() +def quantize_weight_per_tensor_absmax(w, n_bits=8): + # w: (out_features, in_features) + scales = w.abs().max() + q_max = 2 ** (n_bits - 1) - 1 + scales.clamp_(min=1e-5).div_(q_max) + w.div_(scales).round_().mul_(scales) + return w + +class W8A16Linear(nn.Module): + def __init__( + self, + # bit_width, + in_features, + out_features, + bias=True, + quantize_output=False, + ): + super().__init__() + # self.bit_width = bit_width + self.in_features = in_features + self.out_features = out_features + + self.register_buffer( + "weight", + torch.randn( + self.out_features, + self.in_features, + dtype=torch.float16, + requires_grad=False, + ), + ) + if bias: + self.register_buffer( + "bias", + torch.zeros( + (1, self.out_features), dtype=torch.float16, requires_grad=False + ), + ) + else: + self.register_buffer("bias", None) + + def to(self, *args, **kwargs): + super(W8A16Linear, self).to(*args, **kwargs) + self.weight = self.weight.to(*args, **kwargs) + if self.bias is not None: + self.bias = self.bias.to(*args, **kwargs) + return self + + @torch.no_grad() + def forward(self, x): + y = torch.functional.F.linear(x, self.weight, self.bias) + return y + + @staticmethod + def from_float( + bit, module, weight_quant="per_channel", quantize_output=False + ): + assert isinstance(module, torch.nn.Linear) + new_module = W8A16Linear( + # bit, + module.in_features, + module.out_features, + module.bias is not None, + quantize_output=quantize_output, + ) + if weight_quant == "per_channel": + new_module.weight = quantize_weight_per_channel_absmax(module.weight, bit) + elif weight_quant == "per_tensor": + new_module.weight = quantize_weight_per_tensor_absmax(module.weight, bit) + else: + raise ValueError(f"Invalid weight_quant: {weight_quant}") + new_module.weight_quant_name = weight_quant + if module.bias is not None: + new_module.bias = module.bias + return new_module + + def __repr__(self): + return f"W8A16Linear({self.in_features}, {self.out_features}, bias={self.bias is not None}, weight_quant={self.weight_quant_name})" + +def quantize_llama_like( + model, mlp_quant, self_attn_quant, low_bit, weight_quant="per_channel", quantize_bmm_input=False +): + from transformers.models.llama.modeling_llama import ( + LlamaAttention, + LlamaMLP, + ) + + for name, m in model.model.named_modules(): + if isinstance(m, LlamaMLP): + if low_bit == 0: + continue + else: + if name in mlp_quant: + bit = low_bit + print(f'{name} {bit} bit quant ') + else: + if low_bit == 4: + bit = 8 + print(f'{name} {bit} bit quant ') + elif low_bit == 8: + continue + + m.gate_proj = W8A16Linear.from_float( + bit, m.gate_proj, weight_quant=weight_quant + ) + m.up_proj = W8A16Linear.from_float( + bit, m.up_proj, weight_quant=weight_quant + ) + m.down_proj = W8A16Linear.from_float( + bit, m.down_proj, weight_quant=weight_quant + ) + elif isinstance(m, LlamaAttention): + if low_bit == 0: + continue + else: + if name in self_attn_quant: + bit = low_bit + else: + if low_bit == 4: + bit = 8 + elif low_bit == 8: + continue + + m.q_proj = W8A16Linear.from_float( + bit, + m.q_proj, + weight_quant=weight_quant, + quantize_output=quantize_bmm_input, + ) + m.k_proj = W8A16Linear.from_float( + bit, + m.k_proj, + weight_quant=weight_quant, + quantize_output=quantize_bmm_input, + ) + m.v_proj = W8A16Linear.from_float( + bit, + m.v_proj, + weight_quant=weight_quant, + quantize_output=quantize_bmm_input, + ) + m.o_proj = W8A16Linear.from_float( + bit, m.o_proj, weight_quant=weight_quant + ) + + return model \ No newline at end of file diff --git a/LSAQ_CoreCode/main.ipynb b/LSAQ_CoreCode/main.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..3bd73b3e4767e3b11d0286d3fabdce2274427164 --- /dev/null +++ b/LSAQ_CoreCode/main.ipynb @@ -0,0 +1,321 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "import torch\n", + "import torch.nn as nn\n", + "import GPUtil\n", + "from transformers import AutoTokenizer, AutoModelForCausalLM\n", + "import tqdm\n", + "from functools import partial" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Resource Detection" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "gpus = GPUtil.getGPUs()\n", + "free_memory = []\n", + "\n", + "for gpu in gpus:\n", + " free_memory.append(gpu.memoryFree)\n", + "\n", + "memory_sort = sorted(range(len(free_memory)), key=lambda i: free_memory[i])\n", + "\n", + "gpu_id = memory_sort[-1]\n", + "gpu_memory = free_memory[memory_sort[-1]]\n", + "\n", + "print(f'gpu_id:{gpu_id}; gpu_memory:{gpu_memory}')\n", + "\n", + "os.environ[\"CUDA_DEVICE_ORDER\"] = \"PCI_BUS_ID\"\n", + "os.environ[\"CUDA_VISIBLE_DEVICES\"] = str(gpu_id)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Model Selection" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "model_name = \"/data/LLMs/Llama-2-7b-hf\"\n", + "\n", + "tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)\n", + "model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True, torch_dtype=torch.float16, device_map=\"auto\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Layer Importance Detection" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def encode(tok, text, padding=True, truncation=True, max_length=None):\n", + " # 将文本转换为输入 IDs\n", + " input_ids = [tok.bos_id] + tok.encode(text)\n", + "\n", + " # 生成注意力掩码\n", + " attention_mask = [1] * len(input_ids)\n", + "\n", + " # 如果进行了填充，则调整注意力掩码\n", + " if padding:\n", + " padding_length = max_length - len(input_ids)\n", + " attention_mask = [0] * padding_length + attention_mask\n", + " input_ids = [tok.eos_id] * padding_length + input_ids\n", + "\n", + " encoded_input = {\n", + " 'input_ids': input_ids,\n", + " 'attention_mask': attention_mask\n", + " }\n", + " return encoded_input\n", + "\n", + "def batch_encode_plus(tok, texts, max_length=None, return_tensors=None):\n", + " encoded_inputs = []\n", + "\n", + " # 循环处理每个文本\n", + " if max_length is None:\n", + " max_length = -1\n", + " for text in texts:\n", + " # if isinstance(text, list):\n", + " # text = text[0]\n", + " # print(text)\n", + " len_ = len([tok.bos_id] + tok.encode(text))\n", + " if len_ > max_length:\n", + " max_length = len_\n", + " for text in texts:\n", + " # if isinstance(text, list):\n", + " # text = text[0]\n", + " encoded_input = encode(tok, text, max_length = max_length)\n", + " encoded_inputs.append(encoded_input)\n", + "\n", + " # 合并结果\n", + " batch_encoded = {\n", + " 'input_ids': [encoded_input['input_ids'] for encoded_input in encoded_inputs],\n", + " 'attention_mask': [encoded_input['attention_mask'] for encoded_input in encoded_inputs]\n", + " }\n", + "\n", + " batch_encoded = {key: torch.tensor(val) for key, val in batch_encoded.items()}\n", + "\n", + " return batch_encoded" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "tokenizer.bos_token = tokenizer.eos_token\n", + "tokenizer.bos_id = tokenizer.bos_token_id\n", + "tokenizer.eos_id = tokenizer.eos_token_id\n", + "importances = [0 for i in range(len(model.model.layers))] # layer-wise importance scores" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from datasets import load_dataset\n", + "\n", + "dataset = load_dataset(\"wikitext\", \"wikitext-2-raw-v1\", split=\"test\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "MAX_SEQ_LEN = 1024\n", + "batch_size = 1\n", + "dataset_size = 200" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def jaccard_set(list1, list2):\n", + " \"\"\"Define Jaccard Similarity function for two sets\"\"\"\n", + " intersection = len(list(set(list1).intersection(list2)))\n", + " union = (len(list1) + len(list2)) - intersection\n", + " return float(intersection) / union" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "\n", + "k = 20\n", + "\n", + "for i in tqdm.tqdm(range(0, dataset_size, batch_size), total = dataset_size / batch_size):\n", + " \n", + " prompts = dataset['text'][i:i + batch_size]\n", + " max_seq_len = MAX_SEQ_LEN\n", + " stride = 256\n", + " max_gen_len = 0\n", + "\n", + "\n", + " prompt_tokens = batch_encode_plus(\n", + " tokenizer,\n", + " prompts,\n", + " return_tensors='pt'\n", + " )\n", + " input_ids = prompt_tokens['input_ids']\n", + " attn_mask = prompt_tokens['attention_mask']\n", + " max_prompt_len = max(len(t) for t in input_ids)\n", + " all_jac_sim = [0 for i in range(len(model.model.layers))] \n", + " E = model.get_input_embeddings().weight.detach()\n", + " \n", + " # authors use a sliding window of size 1024 with a shift of 256\n", + " for start in range(0, max_prompt_len, stride):\n", + " seq_ids = (attn_mask.sum(dim=-1) > start).nonzero().squeeze()\n", + " seq_ids = seq_ids.unsqueeze(0) if seq_ids.dim() == 0 else seq_ids # ensure 2d\n", + " inputs = input_ids[seq_ids, start:start+max_seq_len]\n", + " attn = attn_mask[seq_ids, start:start+max_seq_len]\n", + "\n", + " if max_gen_len == 0:\n", + " outputs = model(\n", + " input_ids=inputs.to(\"cuda\"),\n", + " attention_mask=attn.to(\"cuda\"),\n", + " output_hidden_states=True,\n", + " )\n", + " else:\n", + " outputs = model.generate(\n", + " input_ids=inputs.to(\"cuda\"),\n", + " attention_mask=attn.to(\"cuda\"),\n", + " max_new_tokens=max_gen_len, \n", + " output_hidden_states=True,\n", + " return_dict_in_generate=True,\n", + " )\n", + "\n", + " hiddens = outputs.hidden_states\n", + "\n", + " for i in range(len(hiddens) - 1):\n", + " in_hidden = hiddens[i][:,-1,:]\n", + " out_hidden = hiddens[i+1][:,-1,:]\n", + "\n", + " in_projs = in_hidden @ E.T\n", + " out_projs = out_hidden @ E.T\n", + "\n", + " in_projs = in_projs.detach().cpu().numpy()\n", + " ot_projs = out_projs.detach().cpu().numpy()\n", + "\n", + " in_ind = np.argsort(-in_projs)\n", + " ot_ind = np.argsort(-ot_projs)\n", + "\n", + " in_topks = [tokenizer.decode(i) for i in in_ind[0][:k]]\n", + " ot_topks = [tokenizer.decode(i) for i in ot_ind[0][:k]]\n", + "\n", + " all_jac_sim[i] += jaccard_set(in_topks, ot_topks)\n", + "\n", + " \n", + " importances = [x + y for x, y in zip(importances, all_jac_sim)]\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import math\n", + "def normalize(lst, range_min=0, range_max=1):\n", + " min_val = min(lst)\n", + " max_val = max(lst)\n", + " normalized = [(range_max - range_min) * (x - min_val) / (max_val - min_val) + range_min for x in lst]\n", + " return normalized\n", + "\n", + "filtered_values = [0 if math.isinf(value) else value for value in importances] \n", + "normalized_lst = normalize(filtered_values)\n", + "\n", + "sorted_indices = sorted(range(len(normalized_lst)), key=lambda i: normalized_lst[i])\n", + "reversed_list = list(reversed(sorted_indices))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Quantize" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from lsaq_quant import quantize_llama_like\n", + "\n", + "num_of_layer2quant = 8\n", + "bit = 8\n", + "\n", + "layer_to_quant = reversed_list[0:num_of_layer2quant]\n", + "\n", + "mlp_quant = [f'layers.{item}.mlp' for item in layer_to_quant]\n", + "self_attn_quant = [f'layers.{item}.self_attn' for item in layer_to_quant]\n", + "\n", + "print(f'quanting ... ')\n", + "model_lsaq = quantize_llama_like(model, mlp_quant, self_attn_quant, bit)\n", + "print(f'quanted')" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "smoothquant", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.8.19" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/README.md b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..e1f08db8d1c83d85e73a71c4785725c8f2c5d600 --- /dev/null +++ b/README.md @@ -0,0 +1 @@ +# quantization \ No newline at end of file diff --git a/__pycache__/zscore.cpython-310.pyc b/__pycache__/zscore.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f019d7e95de416bd885b6b7da640c55b527a5d92 Binary files /dev/null and b/__pycache__/zscore.cpython-310.pyc differ diff --git a/__pycache__/zscore.cpython-311.pyc b/__pycache__/zscore.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..12cbb90eb09cc09b4e41916ccd18e9f629fa672f Binary files /dev/null and b/__pycache__/zscore.cpython-311.pyc differ diff --git a/baselines/Llama-2-7b-hf_alpha_idx_10.json b/baselines/Llama-2-7b-hf_alpha_idx_10.json new file mode 100644 index 0000000000000000000000000000000000000000..e948afb0fcf30e40ab627bbae3f1182ae0dbb07c --- /dev/null +++ b/baselines/Llama-2-7b-hf_alpha_idx_10.json @@ -0,0 +1,34 @@ +[ + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 2, + 2, + 2, + 2, + 2, + 2, + 4, + 2, + 2, + 2, + 2, + 4 +] \ No newline at end of file diff --git a/baselines/Llama-2-7b-hf_alpha_idx_5.json b/baselines/Llama-2-7b-hf_alpha_idx_5.json new file mode 100644 index 0000000000000000000000000000000000000000..41620c390d6fd8ea9eecb80c44fbdf719c6f1b44 --- /dev/null +++ b/baselines/Llama-2-7b-hf_alpha_idx_5.json @@ -0,0 +1,34 @@ +[ + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 2, + 2, + 2, + 4, + 2, + 4, + 4, + 4, + 4, + 4, + 2, + 4 +] \ No newline at end of file diff --git a/baselines/Llama-2-7b-hf_kurtosis_idx_10.json b/baselines/Llama-2-7b-hf_kurtosis_idx_10.json new file mode 100644 index 0000000000000000000000000000000000000000..0ac0e90089a37bb5a0a3427233dceabce7bed3de --- /dev/null +++ b/baselines/Llama-2-7b-hf_kurtosis_idx_10.json @@ -0,0 +1,34 @@ +[ + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 2, + 2, + 4, + 2, + 2, + 2, + 4, + 2, + 4, + 2, + 4, + 2, + 4, + 2, + 2, + 4, + 4, + 4 +] \ No newline at end of file diff --git a/baselines/Llama-2-7b-hf_kurtosis_idx_5.json b/baselines/Llama-2-7b-hf_kurtosis_idx_5.json new file mode 100644 index 0000000000000000000000000000000000000000..d21b7074b69afa377c31b1f8e836769bd03d2259 --- /dev/null +++ b/baselines/Llama-2-7b-hf_kurtosis_idx_5.json @@ -0,0 +1,34 @@ +[ + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 2, + 4, + 4, + 4, + 4, + 2, + 4, + 2, + 4, + 2, + 2, + 4, + 4, + 4 +] \ No newline at end of file diff --git a/baselines/Llama-2-7b-hf_z_idx_10.json b/baselines/Llama-2-7b-hf_z_idx_10.json new file mode 100644 index 0000000000000000000000000000000000000000..f9c210e0e88b56cfc45828d1f26aab3cc4341fb2 --- /dev/null +++ b/baselines/Llama-2-7b-hf_z_idx_10.json @@ -0,0 +1,34 @@ +[ + 4, + 4, + 4, + 2, + 2, + 2, + 4, + 2, + 4, + 2, + 2, + 4, + 2, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 2, + 4, + 4, + 4, + 2, + 4, + 2, + 4, + 4, + 4, + 4 +] \ No newline at end of file diff --git a/baselines/Llama-2-7b-hf_z_idx_5.json b/baselines/Llama-2-7b-hf_z_idx_5.json new file mode 100644 index 0000000000000000000000000000000000000000..93f91c3050675723ff1463c79c824e2734fab004 --- /dev/null +++ b/baselines/Llama-2-7b-hf_z_idx_5.json @@ -0,0 +1,34 @@ +[ + 4, + 4, + 4, + 4, + 4, + 2, + 4, + 4, + 4, + 4, + 4, + 4, + 2, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 2, + 4, + 4, + 4, + 2, + 4, + 2, + 4, + 4, + 4, + 4 +] \ No newline at end of file diff --git a/baselines_1/Llama-2-7b-hf_bi_idx_10.json b/baselines_1/Llama-2-7b-hf_bi_idx_10.json new file mode 100644 index 0000000000000000000000000000000000000000..4548c1bb35b151f690771aadba9ec984b0bfe05a --- /dev/null +++ b/baselines_1/Llama-2-7b-hf_bi_idx_10.json @@ -0,0 +1,34 @@ +[ + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 4, + 4 +] \ No newline at end of file diff --git a/baselines_1/Llama-2-7b-hf_bi_idx_5.json b/baselines_1/Llama-2-7b-hf_bi_idx_5.json new file mode 100644 index 0000000000000000000000000000000000000000..fe2595cad704583b8cb3e90f038f7be61e95904d --- /dev/null +++ b/baselines_1/Llama-2-7b-hf_bi_idx_5.json @@ -0,0 +1,34 @@ +[ + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 2, + 2, + 2, + 2, + 2, + 4, + 4, + 4 +] \ No newline at end of file diff --git a/baselines_1/Llama-2-7b-hf_zd_idx_10.json b/baselines_1/Llama-2-7b-hf_zd_idx_10.json new file mode 100644 index 0000000000000000000000000000000000000000..8c22663c27c654e791172192258b277f329cb4ec --- /dev/null +++ b/baselines_1/Llama-2-7b-hf_zd_idx_10.json @@ -0,0 +1,34 @@ +[ + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 2, + 2, + 4, + 2, + 2, + 2, + 2, + 4, + 2, + 2, + 2, + 2, + 4, + 4, + 4 +] \ No newline at end of file diff --git a/baselines_1/Llama-2-7b-hf_zd_idx_5.json b/baselines_1/Llama-2-7b-hf_zd_idx_5.json new file mode 100644 index 0000000000000000000000000000000000000000..f147c5627c10f8d78e815f510bacaf8d7c3f14f5 --- /dev/null +++ b/baselines_1/Llama-2-7b-hf_zd_idx_5.json @@ -0,0 +1,34 @@ +[ + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 2, + 2, + 4, + 2, + 4, + 2, + 2, + 4, + 4, + 4 +] \ No newline at end of file diff --git a/eval.sh b/eval.sh new file mode 100644 index 0000000000000000000000000000000000000000..d3092ca054b58d1c09b50256c541fa6d9d83bb38 --- /dev/null +++ b/eval.sh @@ -0,0 +1,59 @@ +export HTTP_PROXY=http://sys-proxy-rd-relay.byted.org:8118 +export http_proxy=http://sys-proxy-rd-relay.byted.org:8118 +export https_proxy=http://sys-proxy-rd-relay.byted.org:8118 +export no_proxy="$no_proxy,.byteintl.net" +export HF_ENDPOINT=https://hf-mirror.com + + +cd quantization_metric/ +model_id=/mnt/bn/life-mllm/users/cxr/quantization/models/Qwen/Qwen2.5-14B +model_name=$(basename "$model_id") +cuda_id=4 + +model=/mnt/bn/life-mllm/users/cxr/quantization/models/${model_name}-quantization-layer-mlp +# output_dir=Alpha_values_mlp +# bit_layers_dir=/mnt/bn/life-mllm/users/cxr/quantization/quantization_metric/bit_layers +# result_dir=/mnt/bn/life-mllm/users/cxr/quantization/lm-evaluation-harness/results + + + +start=$(date +%s.%N) +# rm -rf $model + +modes=("mlp") +for mode in ${modes[@]}; do + for idx in {32..47}; do + echo $mode $idx + cd ../quantization_metric + python -u main.py --bit_layer_idx $idx --save_dir ${model} --mode $mode --model_id $model_id --cuda_id $cuda_id + cd ../lm-evaluation-harness + bash run_scripts/eval_base_qwen2_5_14b.sh ${model} ${mode}_${idx} + rm -rf ${model} + end=$(date +%s.%N) + runtime=$(awk "BEGIN {print $end - $start}") + echo "Execution time: $runtime seconds" + + done +done + + +model=/mnt/bn/life-mllm/users/cxr/quantization/models/${model_name}-quantization-layer + +start=$(date +%s.%N) +# rm -rf $model + +modes=("self_attn") +for mode in ${modes[@]}; do + for idx in {32..47}; do + echo $mode $idx + cd ../quantization_metric + python -u main.py --bit_layer_idx $idx --save_dir ${model} --mode $mode --model_id $model_id --cuda_id $cuda_id + cd ../lm-evaluation-harness + bash run_scripts/eval_base_qwen2_5_14b.sh ${model} ${mode}_${idx} + rm -rf ${model} + end=$(date +%s.%N) + runtime=$(awk "BEGIN {print $end - $start}") + echo "Execution time: $runtime seconds" + + done +done \ No newline at end of file diff --git a/eval_coherence.sh b/eval_coherence.sh new file mode 100644 index 0000000000000000000000000000000000000000..ac39c2c8b7d08e1abdda7709f1bf98d5085f2798 --- /dev/null +++ b/eval_coherence.sh @@ -0,0 +1,29 @@ +export HTTP_PROXY=http://sys-proxy-rd-relay.byted.org:8118 +export http_proxy=http://sys-proxy-rd-relay.byted.org:8118 +export https_proxy=http://sys-proxy-rd-relay.byted.org:8118 +export no_proxy="$no_proxy,.byteintl.net" +export HF_ENDPOINT=https://hf-mirror.com + +cd quantization_metric/ +model=../models/patch/Llama-2-7b-hf-quantization +# output_dir=Alpha_values_mlp +tasks=piqa,winogrande,arc_easy,arc_challenge,hellaswag,boolq +# bit_layers_dir=/mnt/bn/life-mllm/users/cxr/quantization/quantization_metric/bit_layers +# result_dir=/mnt/bn/life-mllm/users/cxr/quantization/lm-evaluation-harness/results + + +start=$(date +%s.%N) +# rm -rf $model + +file=/mnt/bn/life-mllm/users/cxr/quantization/quantization_metric/metrics/coherence/coherence_self_attn_Llama-2-7b-hf.json +echo "$file" +cd ../quantization_metric +configure_id=$(basename $file .json) +python -u main_low.py --bit_layers $file --save_dir ${model} --k 5 --reverse False +cd ../lm-evaluation-harness +bash run_scripts/eval.sh ${configure_id} ${model} ${tasks} +rm -rf ${model} +end=$(date +%s.%N) +runtime=$(awk "BEGIN {print $end - $start}") +echo "Execution time: $runtime seconds" + diff --git a/eval_fg.sh b/eval_fg.sh new file mode 100644 index 0000000000000000000000000000000000000000..912ddfb5f729e089e255fb0bc438fb1caf79ea4a --- /dev/null +++ b/eval_fg.sh @@ -0,0 +1,89 @@ +export HTTP_PROXY=http://sys-proxy-rd-relay.byted.org:8118 +export http_proxy=http://sys-proxy-rd-relay.byted.org:8118 +export https_proxy=http://sys-proxy-rd-relay.byted.org:8118 +export no_proxy="$no_proxy,.byteintl.net" +export HF_ENDPOINT=https://hf-mirror.com + + +cd quantization_metric/ +model_id=/mnt/bn/life-mllm/users/cxr/quantization/models/Qwen/Qwen2.5-7B +model_name=$(basename "$model_id") + +model=/mnt/bn/life-mllm/users/cxr/quantization/models/${model_name}-quantization-fg +# output_dir=Alpha_values_mlp +# bit_layers_dir=/mnt/bn/life-mllm/users/cxr/quantization/quantization_metric/bit_layers +# result_dir=/mnt/bn/life-mllm/users/cxr/quantization/lm-evaluation-harness/results + +start=$(date +%s.%N) +# rm -rf $model +cd ../quantization_metric +# fg1 +# self_attn_layer_to_quant="4 1 2 8 23" +# mlp_layer_to_quant="27 16 19 17 25" + + +# save_fg=fg2 +# self_attn_layer_to_quant="23 22 25 24 26" +# mlp_layer_to_quant="27 16 19 17 25" + + +# save_fg=fg3 +# self_attn_layer_to_quant="23 22 25 24 26" +# mlp_layer_to_quant="27 16 19" + + + +# save_fg=fg4 +# self_attn_layer_to_quant="23 22 25 24 26" +# mlp_layer_to_quant="27" + + +# save_fg=fg5 +# self_attn_layer_to_quant="27 16 19 17 25" +# mlp_layer_to_quant="27 16 19 17 25" + +# save_fg=baseline_BI +# self_attn_layer_to_quant="16 17 15 14 13" +# mlp_layer_to_quant="16 17 15 14 13" + + +save_fg=f6 +self_attn_layer_to_quant="4 1 2 8 23 22 25 5 24 7 26 6 20 12 19 17 21 11 10 9 18" +mlp_layer_to_quant="27 16 19" + +python -u main_fg.py --cuda_id 6 --save_dir ${model} --model_id $model_id --self_attn_layer_to_quant "${self_attn_layer_to_quant}" --mlp_layer_to_quant "${mlp_layer_to_quant}" +cd ../lm-evaluation-harness +bash run_scripts/eval_base_fg.sh ${model} ${save_fg} +rm -rf ${model} +end=$(date +%s.%N) +runtime=$(awk "BEGIN {print $end - $start}") +echo "Execution time: $runtime seconds" + + + +model_id=/mnt/bn/life-mllm/users/cxr/quantization/models/meta-llama/Llama-3.1-8B +model_name=$(basename "$model_id") + +model=/mnt/bn/life-mllm/users/cxr/quantization/models/${model_name}-quantization-fg +# output_dir=Alpha_values_mlp +# bit_layers_dir=/mnt/bn/life-mllm/users/cxr/quantization/quantization_metric/bit_layers +# result_dir=/mnt/bn/life-mllm/users/cxr/quantization/lm-evaluation-harness/results + +start=$(date +%s.%N) +# rm -rf $model +cd ../quantization_metric +# save_fg=baseline_BI +# self_attn_layer_to_quant="24 25 23 26 27" +# mlp_layer_to_quant="24 25 23 26 27" +save_fg=fg6 +self_attn_layer_to_quant="29 23 24 30 18 28 26 20 16 27 25 17 19 21" +mlp_layer_to_quant="26 20 22" + +python -u main_fg.py --cuda_id 6 --save_dir ${model} --model_id $model_id --self_attn_layer_to_quant "${self_attn_layer_to_quant}" --mlp_layer_to_quant "${mlp_layer_to_quant}" +cd ../lm-evaluation-harness +bash run_scripts/eval_base_fg.sh ${model} ${save_fg} +rm -rf ${model} +end=$(date +%s.%N) +runtime=$(awk "BEGIN {print $end - $start}") +echo "Execution time: $runtime seconds" + diff --git a/eval_hd.sh b/eval_hd.sh new file mode 100644 index 0000000000000000000000000000000000000000..867983c961a3eada02d72f6b5c202500ffde7718 --- /dev/null +++ b/eval_hd.sh @@ -0,0 +1,28 @@ +export HTTP_PROXY=http://sys-proxy-rd-relay.byted.org:8118 +export http_proxy=http://sys-proxy-rd-relay.byted.org:8118 +export https_proxy=http://sys-proxy-rd-relay.byted.org:8118 +export no_proxy="$no_proxy,.byteintl.net" +export HF_ENDPOINT=https://hf-mirror.com + +cd quantization_metric/ +model=../models/patch/Llama-2-7b-hf-quantization +# output_dir=Alpha_values_mlp +tasks=piqa,winogrande,arc_easy,arc_challenge,hellaswag,boolq +# bit_layers_dir=/mnt/bn/life-mllm/users/cxr/quantization/quantization_metric/bit_layers +# result_dir=/mnt/bn/life-mllm/users/cxr/quantization/lm-evaluation-harness/results + + +start=$(date +%s.%N) +# rm -rf $model + +file=/mnt/bn/life-mllm/users/cxr/quantization/quantization_metric/metrics/head_diversity/head_diversity_self_attn_Llama-2-7b-hf.json +echo "$file" +cd ../quantization_metric +configure_id=$(basename $file .json) +python -u main_low.py --bit_layers $file --save_dir ${model} --k 5 --reverse False +cd ../lm-evaluation-harness +bash run_scripts/eval.sh ${configure_id} ${model} ${tasks} +rm -rf ${model} +end=$(date +%s.%N) +runtime=$(awk "BEGIN {print $end - $start}") +echo "Execution time: $runtime seconds" diff --git a/eval_layer_llama.sh b/eval_layer_llama.sh new file mode 100644 index 0000000000000000000000000000000000000000..14f3d9475be49bb07e0edc10610138234f4d4cf0 --- /dev/null +++ b/eval_layer_llama.sh @@ -0,0 +1,39 @@ +export HTTP_PROXY=http://sys-proxy-rd-relay.byted.org:8118 +export http_proxy=http://sys-proxy-rd-relay.byted.org:8118 +export https_proxy=http://sys-proxy-rd-relay.byted.org:8118 +export no_proxy="$no_proxy,.byteintl.net" +export HF_ENDPOINT=https://hf-mirror.com + + +cd quantization_metric/ +cuda_id=0 +model_id="/mnt/bn/life-mllm/users/cxr/quantization/models/meta-llama/Llama-3.1-8B" +model_name=$(basename "$model_id") + +model=/mnt/bn/life-mllm/users/cxr/quantization/models/${model_name}-quantization-layer +# output_dir=Alpha_values_mlp +# bit_layers_dir=/mnt/bn/life-mllm/users/cxr/quantization/quantization_metric/bit_layers +# result_dir=/mnt/bn/life-mllm/users/cxr/quantization/lm-evaluation-harness/results + + + +start=$(date +%s.%N) +# rm -rf $model + +modes=("self_attn" "mlp") +for mode in ${modes[@]}; do + for idx in {-1..31}; do + echo $mode $idx + cd ../quantization_metric + python -u main.py --bit_layer_idx $idx --save_dir ${model} --mode $mode --model_id $model_id --cuda_id $cuda_id + cd ../lm-evaluation-harness + bash run_scripts/eval_base_qwen2_5_7b.sh ${model} ${mode}_${idx} $cuda_id + rm -rf ${model} + end=$(date +%s.%N) + runtime=$(awk "BEGIN {print $end - $start}") + echo "Execution time: $runtime seconds" + + done +done + + diff --git a/eval_layer_qwen.sh b/eval_layer_qwen.sh new file mode 100644 index 0000000000000000000000000000000000000000..cf4d4139cc49e1fecac0471838d5b907cc157092 --- /dev/null +++ b/eval_layer_qwen.sh @@ -0,0 +1,39 @@ +export HTTP_PROXY=http://sys-proxy-rd-relay.byted.org:8118 +export http_proxy=http://sys-proxy-rd-relay.byted.org:8118 +export https_proxy=http://sys-proxy-rd-relay.byted.org:8118 +export no_proxy="$no_proxy,.byteintl.net" +export HF_ENDPOINT=https://hf-mirror.com + + +cd quantization_metric/ +cuda_id=1 +model_id=/mnt/bn/life-mllm/users/cxr/quantization/models/Qwen/Qwen2.5-7B +model_name=$(basename "$model_id") + +model=/mnt/bn/life-mllm/users/cxr/quantization/models/${model_name}-quantization-layer +# output_dir=Alpha_values_mlp +# bit_layers_dir=/mnt/bn/life-mllm/users/cxr/quantization/quantization_metric/bit_layers +# result_dir=/mnt/bn/life-mllm/users/cxr/quantization/lm-evaluation-harness/results + + + +start=$(date +%s.%N) +# rm -rf $model + +modes=("self_attn" "mlp") +for mode in ${modes[@]}; do + for idx in {-1..27}; do + echo $mode $idx + cd ../quantization_metric + python -u main.py --bit_layer_idx $idx --save_dir ${model} --mode $mode --model_id $model_id --cuda_id $cuda_id + cd ../lm-evaluation-harness + bash run_scripts/eval_base_qwen2_5_7b.sh ${model} ${mode}_${idx} $cuda_id + rm -rf ${model} + end=$(date +%s.%N) + runtime=$(awk "BEGIN {print $end - $start}") + echo "Execution time: $runtime seconds" + + done +done + + diff --git a/eval_zd.sh b/eval_zd.sh new file mode 100644 index 0000000000000000000000000000000000000000..cb568fb9e70363bc57951dceed1df207116f255a --- /dev/null +++ b/eval_zd.sh @@ -0,0 +1,29 @@ +export HTTP_PROXY=http://sys-proxy-rd-relay.byted.org:8118 +export http_proxy=http://sys-proxy-rd-relay.byted.org:8118 +export https_proxy=http://sys-proxy-rd-relay.byted.org:8118 +export no_proxy="$no_proxy,.byteintl.net" +export HF_ENDPOINT=https://hf-mirror.com + +cd quantization_metric/ +model=../models/patch/Llama-2-7b-hf-quantization-zd +# output_dir=Alpha_values_mlp +tasks=piqa,winogrande,arc_easy,arc_challenge,hellaswag,boolq +# bit_layers_dir=/mnt/bn/life-mllm/users/cxr/quantization/quantization_metric/bit_layers +# result_dir=/mnt/bn/life-mllm/users/cxr/quantization/lm-evaluation-harness/results + + +start=$(date +%s.%N) +# rm -rf $model + +# file=/mnt/bn/life-mllm/users/cxr/quantization/quantization_metric/metrics/ZD/ZD_mlp_Llama-2-7b-hf.json +file=/mnt/bn/life-mllm/users/cxr/quantization/quantization_metric/metrics/ZD/ZD_self_attn_Llama-2-7b-hf.json +echo "$file" +cd ../quantization_metric +configure_id=$(basename $file .json) +python -u main_low.py --bit_layers $file --save_dir ${model} --k 5 +cd ../lm-evaluation-harness +bash run_scripts/eval.sh ${configure_id} ${model} ${tasks} +rm -rf ${model} +end=$(date +%s.%N) +runtime=$(awk "BEGIN {print $end - $start}") +echo "Execution time: $runtime seconds" diff --git a/inference.py b/inference.py new file mode 100644 index 0000000000000000000000000000000000000000..b5b20419c0174551c3a9857dd4f136ace4ba60a5 --- /dev/null +++ b/inference.py @@ -0,0 +1,19 @@ + +a= {"results": { + "arc_easy": { + "alias": "arc_easy", + "acc,none": 0.6902356902356902, + "acc_stderr,none": 0.00948817285190372, + "acc_norm,none": 0.6422558922558923, + "acc_norm_stderr,none": 0.00983577275734336 + }, + "arc_easy": { + "alias": "arc_easy", + "acc,none": 0.6902356902356902, + "acc_stderr,none": 0.00948817285190372, + "acc_norm,none": 0.6422558922558923, + "acc_norm_stderr,none": 0.00983577275734336 + } + } +} +print(len(a['results'])) \ No newline at end of file diff --git a/layerwise-awq.py b/layerwise-awq.py new file mode 100644 index 0000000000000000000000000000000000000000..280dd4591eb63c84b3b2ff6f5c4b3f893ae589ad --- /dev/null +++ b/layerwise-awq.py @@ -0,0 +1,322 @@ +# -*- encoding:utf-8 -*- +@torch.no_grad() +def run_awq( + model, + enc, + w_bit, + q_config, + n_samples=512, + seqlen=512, + auto_scale=True, + mse_range=True, + calib_data="pileval", # data for calibration + skip_first: int = 0, # number of initial layers to keep in full precision + first_n: int = 0, # number of initial layers to apply first quant + w_bit_first: int | None = None, + w_bit_rest: int | None = None, + # --- mixed-precision strategy -------------------------------------------------- + strategy: str = "layer", # "layer" (default): original solve layer-by-layer; "auto": structured mixed-precision + m_auto: int | None = None, # number of high-bit layers when strategy == "auto"; defaults to 25% of L + hi_bit: int = 4, + lo_bit: int = 2, + alpha: float = 1 / 3, + beta: float = 1 / 3, + gamma: float = 1 / 3, + k_energy: int = 32, + metrics_csv: str | None = None, # optional explicit path to metrics CSV (delta_ppl,erank_diff,topk_energy_diff) +): + from ..utils.calib_data import get_calib_dataset + from ..utils.module import append_str_prefix, get_op_name + + if "bigcode" in str(model.__class__).lower(): + # otherwise attention_mask will always be on cpu. + model.transformer.bias = model.transformer.bias.to("cuda") + + layers = get_blocks(model) + + samples = get_calib_dataset( + data=calib_data, tokenizer=enc, n_samples=n_samples, block_size=seqlen + ) + samples = torch.cat(samples, dim=0) + + inps = [] + layer_kwargs = {} + + layers[0] = layers[0].cuda() + move_embed(model, "cuda") + + # get input and kwargs to layer 0 + # with_kwargs is only supported in PyTorch 2.0 + # use this Catcher hack for now + class Catcher(nn.Module): + def __init__(self, module): + super().__init__() + self.module = module + + def forward(self, inp, **kwargs): + inps.append(inp) + layer_kwargs.update(kwargs) + raise ValueError # early exit to break later inference + + # patch layer 0 to catch input and kwargs + layers[0] = Catcher(layers[0]) + try: + if model.__class__.__name__ == "LlavaLlamaModel": + model.llm(samples.to(next(model.parameters()).device)) + elif model.__class__.__name__ == "InternVL3": + model.language_model(samples.to(next(model.parameters()).device)) + else: + model(samples.to(next(model.parameters()).device)) + except ValueError: # work with early exit + pass + del samples + layers[0] = layers[0].module # restore + inps = inps[0] + + layers[0] = layers[0].cpu() + move_embed(model, "cpu") + + gc.collect() + torch.cuda.empty_cache() + + awq_results = { + "scale": [], + "clip": [], + } + + # --------------------------------------------------------------------------- + # Determine per-layer bit-widths according to the requested *strategy* + # --------------------------------------------------------------------------- + + if strategy.lower() == "auto": + # ------------------------------------------------------------------- + # Use qpRANK pre-computed diagnostics to decide per-layer precision. + # Users may place the JSON files (drop_layer_ppl.json, diff_erank_values.json) + # under the project root (default path) or supply env QPRANK_METRICS_DIR. + # ------------------------------------------------------------------- + + import json, os, math, csv + + def _load_metrics_from_csv(csv_path: str): + """Return delta_ppl, erank_diff, topk_energy_diff lists from a csv file.""" + delta_ppl, erank, topk = [], [], [] + with open(csv_path, "r", encoding="utf-8") as f: + reader = csv.DictReader(f) + for row in reader: + delta_ppl.append(float(row.get("delta_ppl", 0))) + erank.append(abs(float(row.get("erank_diff", 0)))) + topk_val = row.get("topk_energy_diff") + if topk_val is not None and topk_val != "": + topk.append(float(topk_val)) + # Ensure all same length + assert len(delta_ppl) == len(erank), "CSV length mismatch" + if len(topk) != len(delta_ppl): + topk = [0.0] * len(delta_ppl) + return delta_ppl, erank, topk + + delta_ppl: List[float] + delta_r: List[float] + delta_e: List[float] + + # Priority 1: explicit CSV path + if metrics_csv is not None and os.path.isfile(metrics_csv): + delta_ppl, delta_r, delta_e = _load_metrics_from_csv(metrics_csv) + else: + # Priority 2: auto-detect inside QPRANK directory structure + base_dir = os.getenv("QPRANK_METRICS_DIR", os.path.expanduser("~/qpRANK/src")) + + # Derive a crude model identifier from config + cfg_name = getattr(model, "config", None) + model_id = ( + getattr(cfg_name, "_name_or_path", "model").replace("/", "_") + if cfg_name is not None + else "model" + ) + + # Traverse to find a metrics_long.csv matching pattern + candidate_csv = None + for root, dirs, files in os.walk(base_dir): + if "metrics_long.csv" in files and model_id in root: + candidate_csv = os.path.join(root, "metrics_long.csv") + break + + if candidate_csv and os.path.isfile(candidate_csv): + delta_ppl, delta_r, delta_e = _load_metrics_from_csv(candidate_csv) + else: + # Fallback to old JSON files (legacy) + metrics_dir = os.getenv("QPRANK_METRICS_DIR", os.path.expanduser("~/qpRANK")) + ppl_path = os.path.join(metrics_dir, "drop_layer_ppl.json") + erank_path = os.path.join(metrics_dir, "diff_erank_values.json") + + if not (os.path.isfile(ppl_path) and os.path.isfile(erank_path)): + raise FileNotFoundError( + "Cannot locate per-layer metric files for auto strategy. Provide metrics_csv path or set QPRANK_METRICS_DIR appropriately." + ) + + delta_ppl = json.load(open(ppl_path, "r"))["delta_ppl"] + erank_json = json.load(open(erank_path, "r")) + + keys = [k for k in ("q", "k", "v") if k in erank_json] + delta_r = [ + sum(erank_json[k][i] for k in keys) / len(keys) + for i in range(len(delta_ppl)) + ] + + delta_e = erank_json.get("topk_energy_diff", [0.0] * len(delta_ppl)) + #! layer 的数量 + L_total = len(delta_ppl) + + # Normalise + def _norm(arr): + m = max(arr) if max(arr) > 0 else 1.0 + return [x / m for x in arr] + + ppl_hat = _norm(delta_ppl) + r_hat = _norm(delta_r) + e_hat = _norm(delta_e) + + scores = [ + alpha * ppl_hat[i] + beta * r_hat[i] + gamma * e_hat[i] + for i in range(L_total) + ] + + #! 1/4 的 layer + if m_auto is None: + m_auto = max(1, L_total // 4) + + idx_sorted = sorted(range(L_total), key=lambda i: scores[i], reverse=True) + #! 前 1/4 的 layer 用 high bit, 其他的用 low bit + hi_set = set(idx_sorted[:m_auto]) + + #! 每个 layer 的 bit 数量的分配 + #! 我们也是在这边修改成得到我们的 layer 分配就好了 + bits_per_layer = [hi_bit if i in hi_set else lo_bit for i in range(L_total)] + + # ---- verbose print & log ---- + try: + import logging + _logger = logging.getLogger(__name__) + except ImportError: + _logger = None + + print("[AUTO] Per-layer bit-width allocation (index:bit):") + mapping_str = ", ".join(f"{idx}:{bits_per_layer[idx]}b" for idx in range(L_total)) + print(mapping_str) + + if _logger is not None: + _logger.info("AUTO bit-width allocation: " + mapping_str) + + print(f"[AUTO] Layers @ {hi_bit}-bit: {sorted(list(hi_set))}") + print(f"[AUTO] Layers @ {lo_bit}-bit: {sorted([i for i in range(L_total) if i not in hi_set])}") + + if _logger is not None: + _logger.info(f"Layers_{hi_bit}bit: {sorted(list(hi_set))}") + _logger.info(f"Layers_{lo_bit}bit: {[i for i in range(L_total) if i not in hi_set]}") + + else: + # Fallback to original scheme (uniform or head/tail mixed precision). + bits_per_layer = None # will be decided on the fly as before + + # solve layer by layer + for i in tqdm.tqdm(range(len(layers)), desc="Running AWQ..."): + # print(f"Layer {i} of {len(layers)-1}") + layer = layers[i] + + # Flag: whether to apply quantization to this layer + #! 他们也指定了超参数从第几层开始量化 + quantize_this = i >= skip_first + + # Determine bit-width for this layer + if strategy.lower() == "auto" and bits_per_layer is not None: + current_w_bit = bits_per_layer[i] + if i == 0: + # show a brief summary once for user awareness + print( + f"[AUTO] Using structured mixed-precision: {sum(b == hi_bit for b in bits_per_layer)} layers @ {hi_bit}-bit, {sum(b == lo_bit for b in bits_per_layer)} layers @ {lo_bit}-bit." + ) + else: + # original rule-based selection + if i < first_n: + current_w_bit = w_bit_first if w_bit_first is not None else w_bit + print( + f"Layer {i} is quantizing with {current_w_bit} bits. (when this sentence isnt printed, it is quantizing with {w_bit_rest} bits)" + ) + else: + current_w_bit = w_bit_rest if w_bit_rest is not None else w_bit + + + #! 从这边往后就和原来的代码一样 + layer = layer.cuda() + named_linears = get_named_linears(layer) + + # firstly, get input features of all linear layers + def cache_input_hook(m, x, y, name, feat_dict): + x = x[0] + x = x.detach().cpu() + feat_dict[name].append(x) + + input_feat = defaultdict(list) + handles = [] + for name in named_linears: + handles.append( + named_linears[name].register_forward_hook( + functools.partial(cache_input_hook, name=name, feat_dict=input_feat) + ) + ) + inps = inps.to(next(layer.parameters()).device) # in case multi-gpu + # get output as next layer's input + inps = layer(inps, **layer_kwargs)[0] + for h in handles: + h.remove() + # now solve for scaling and clipping + input_feat = {k: torch.cat(v, dim=0) for k, v in input_feat.items()} + + # Clear GPU memory + torch.cuda.empty_cache() + + if ( + auto_scale + ): # if it applies, we should also modify the input_feat with scales + scales_list = auto_scale_block( + layer, + layer_kwargs, + w_bit=current_w_bit, #! 改成 current_w_bit 就可以 + q_config=q_config, + input_feat=input_feat, + ) + # apply_scale(layer, scales_list, input_feat_dict=input_feat) + apply_scale(layers[i], scales_list, input_feat_dict=input_feat) + # append prefix to make names global + awq_results["scale"] += append_str_prefix( + scales_list, get_op_name(model, layer) + "." + ) + + # Clear GPU memory + torch.cuda.empty_cache() + # for line in torch.cuda.memory_summary().splitlines(): + # if "Allocated" in line: + # print(line) + + if mse_range: + clip_list = auto_clip_block( + layer, + w_bit=current_w_bit, #! 改成 current_w_bit 就可以 + q_config=q_config, + input_feat=input_feat, + ) + apply_clip(layer, clip_list) + # append prefix to make names global + awq_results["clip"] += append_str_prefix( + clip_list, get_op_name(model, layer) + "." + ) + + layer = layer.cpu() + # Haotian: check activation replacement + del input_feat + gc.collect() + torch.cuda.empty_cache() + # for line in torch.cuda.memory_summary().splitlines(): + # if "Allocated" in line: + # print(line) + + return awq_results \ No newline at end of file diff --git a/llm-awq/.gitignore b/llm-awq/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..6417a78a68f4d3e0b5cc6728592630d0f8a50c01 --- /dev/null +++ b/llm-awq/.gitignore @@ -0,0 +1,173 @@ +.DS_Store + +data/ +checkpoints +demo_images +serve_images +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +*.pyc +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ +cover/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +.pybuilder/ +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +# For a library or package, you might want to ignore these files since the code is +# intended to run in multiple environments; otherwise, check them in: +# .python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# poetry +# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control +#poetry.lock + +# pdm +# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. +#pdm.lock +# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it +# in version control. +# https://pdm.fming.dev/#use-with-ide +.pdm.toml + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +# Cython debug symbols +cython_debug/ + +# PyCharm +# JetBrains specific template is maintained in a separate JetBrains.gitignore that can +# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore +# and can be added to the global gitignore or merged into this file. For a more nuclear +# option (not recommended) you can uncomment the following to ignore the entire idea folder. +#.idea/ + +*.pt +**/*.pt +**/*.pyc +*.json +__pycache__ diff --git a/llm-awq/LICENSE b/llm-awq/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..aca327a505563c97f6ce15cbb88098e8e72f3965 --- /dev/null +++ b/llm-awq/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2023 MIT HAN Lab + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/llm-awq/README.md b/llm-awq/README.md new file mode 100644 index 0000000000000000000000000000000000000000..2d8fff5ef785f1697dde3744c276d90bae14e5cb --- /dev/null +++ b/llm-awq/README.md @@ -0,0 +1,292 @@ +# AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration +[[Paper](https://arxiv.org/abs/2306.00978)][[Website](https://hanlab.mit.edu/projects/awq)] + +**Efficient and accurate** low-bit weight quantization (INT3/4) for LLMs, supporting **instruction-tuned** models and **multi-modal** LMs. + +![overview](figures/overview.png) + +The current release supports: + +- AWQ search for accurate quantization. +- Pre-computed AWQ model zoo for LLMs (Llama-1/2/3, OPT, CodeLlama, StarCoder, Vicuna, VILA, LLaVA; load to generate quantized weights). +- Memory-efficient 4-bit Linear in PyTorch. +- Efficient CUDA kernel implementation for fast inference (support context and decoding stage). +- Examples on 4-bit inference of an instruction-tuned model (Vicuna) and **multi-modal LM** (VILA). +- Chunk prefilling for faster prefilling in multi-round Q&A setting. +- State-of-the-art prefilling speed of LLMs/VLMs on edge devices: [TinyChat 2.0](./tinychat). + +**Thanks to AWQ, TinyChat can deliver more efficient responses with LLM/VLM chatbots through 4-bit inference.** + +* TinyChat with LLaMA-3-8b on RTX 4090 (2.7x faster than FP16): + +![TinyChat with LLaMA-3-8b on RTX 4090: W4A16 is 2.7x faster than FP16](./tinychat/figures/4090_example_new.gif) + +* TinyChat with LLaMA-3-8b on Jetson Orin (2.9x faster than FP16): + +![TinyChat with LLaMA-3-8b on Jetson Orin: W4A16 is 2.9x faster than FP16](./tinychat/figures/orin_example_new.gif) + + +**TinyChat also supports inference with vision language models (e.g., VILA, LLaVA). In the following examples, W4A16 quantized models from VILA family are launched with TinyChat.** + +* TinyChat with NVILA-8B on RTX 4090 (single-image inputs): + +![TinyChat with NVILA on 4090 single image](./tinychat/figures/4090_nvila_single.gif) + +* TinyChat with NVILA-8B on RTX 4090 (multi-image inputs): + +![TinyChat with NVILA on 4090 multiple images](./tinychat/figures/4090_nvila_multi.gif) + + + +* TinyChat with video reasoning: + +https://github.com/user-attachments/assets/b68a7a0d-5175-4030-985b-5ae0ae94f874 + +**Prompt:** What might be the next step according to the video? + +**Answer:** The next step in the video could be to place the shaped dough onto a baking sheet and let it rise before baking. + +**Online demo:** https://vila.hanlab.ai + +Check out [TinyChat](tinychat), which offers a turn-key solution for **on-device inference** of LLMs and VLMs on **resource-constrained edge platforms**. With TinyChat, it is now possible to efficiently run **large** models on **small** and **low-power** devices even without Internet connection! + + +## News +- [2025/04] 🔥 AWQ now supports DeepSeek-R1-Distilled models. Try our example [here](https://github.com/mit-han-lab/llm-awq/blob/main/scripts/DeepSeek_R1_Distill_example.sh)! +- [2025/02] AWQ now supports BF16 precision. See example [here](https://github.com/mit-han-lab/llm-awq/blob/main/scripts/qwen_example.sh). +- [2024/10] 🔥⚡ Explore advancements in [TinyChat 2.0](./tinychat), the latest version with significant advancements in prefilling speed of Edge LLMs and VLMs, **1.5-1.7x** faster than the previous version of TinyChat. Please refer to the [README](./tinychat/README.md) and [blog](https://hanlab.mit.edu/blog/tinychat20) for more details. +- [2024/05] 🏆 AWQ receives the **Best Paper Award** at **MLSys 2024**. 🎉 +- [2024/05] 🔥 The **VILA-1.5** model family which features **video understanding** is now supported in AWQ and TinyChat. Check out out online demo powered by TinyChat [here](https://vila.hanlab.ai). Example is [here](scripts/vila15_example.sh). +- [2024/05] 🔥 [AMD](https://community.amd.com/t5/ai/reduce-memory-footprint-and-improve-performance-running-llms-on/ba-p/686157) adopts AWQ to improve LLM serving efficiency. +- [2024/04] 🔥 We released AWQ and TinyChat support for The **Llama-3** model family! Check out our example [here](scripts/llama3_example.sh). +- [2024/02] 🔥 AWQ has been accepted to **MLSys 2024**! +- [2024/02] 🔥 We supported [VILA Vision Languague Models](https://arxiv.org/abs/2312.07533) in AWQ & TinyChat! Check our latest demos with multi-image inputs! +- [2024/02] 🔥 We released new version of quantized GEMM/GEMV kernels in [**TinyChat**](tinychat), leading to **38 tokens/second** inference speed on NVIDIA Jetson Orin! +- [2024/01] 🔥 AWQ has been integrated by [Google Vertex AI](https://console.cloud.google.com/vertex-ai/publishers/meta/model-garden/llama-2-quantized)! +- [2023/11] 🔥 AWQ has been integrated by [Amazon Sagemaker Containers](https://aws.amazon.com/blogs/machine-learning/boost-inference-performance-for-llms-with-new-amazon-sagemaker-containers/)! +- [2023/11] 🔥 We added AWQ support and pre-computed search results for CodeLlama, StarCoder, StableCode models. Checkout our model zoo [here](https://huggingface.co/datasets/mit-han-lab/awq-model-zoo)! +- [2023/11] 🔥 AWQ is now integrated natively in Hugging Face transformers through `from_pretrained`. You can either load quantized models from the Hub or your own HF quantized models. +- [2023/10] AWQ is integrated into NVIDIA [TensorRT-LLM](https://github.com/NVIDIA/TensorRT-LLM/) +- [2023/09] AWQ is integrated into [Intel Neural Compressor](https://github.com/intel/neural-compressor), [FastChat](https://github.com/lm-sys/FastChat/blob/main/docs/awq.md), [vLLM](https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/quantization/awq.py), [HuggingFace TGI](https://github.com/huggingface/text-generation-inference/pull/1054), and [LMDeploy](https://github.com/InternLM/lmdeploy). +- [2023/09] ⚡ Check out our latest [**TinyChat**](tinychat), which is ~2x faster than the first release on Orin! +- [2023/09] ⚡ Check out [**AutoAWQ**](https://github.com/casper-hansen/AutoAWQ), a third-party implementation to make AWQ easier to expand to new models, improve inference speed, and integrate into Huggingface. +- [2023/07] 🔥 We released **TinyChat**, an efficient and lightweight chatbot interface based on AWQ. TinyChat enables efficient LLM inference on both cloud and edge GPUs. Llama-2-chat models are supported! Check out our implementation [here](tinychat). +- [2023/07] 🔥 We added AWQ support and pre-computed search results for Llama-2 models (7B & 13B). Checkout our model zoo [here](https://huggingface.co/datasets/mit-han-lab/awq-model-zoo)! +- [2023/07] We extended the support for more LLM models including MPT, Falcon, and BLOOM. + +## Contents + +- [AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration](#awq-activation-aware-weight-quantization-for-llm-compression-and-acceleration) + - [News](#news) + - [Contents](#contents) + - [Helpful Links](#helpful-links) + - [Install](#install) + - [AWQ Model Zoo](#awq-model-zoo) + - [Examples](#examples) + - [Usage](#usage) + - [Results on Visual Language Models](#results-on-visual-language-models) + - [Reference](#reference) + - [Related Projects](#related-projects) + +## Helpful Links + +- [VILA online demo](vila.hanlab.ai): Visual Language Models efficiently supported by AWQ & TinyChat. +- [LLM on the Edge](https://github.com/mit-han-lab/llm-awq/tree/nv_laptop?tab=readme-ov-file#install): AWQ and TinyChat support edge GPUs such as NVIDIA Jetson Orin. +- [VLMs on Laptop](https://github.com/mit-han-lab/llm-awq/tree/nv_laptop?tab=readme-ov-file#run-vila-on-laptop): Follow the instructions to deploy VLMs on NVIDIA Laptops with TinyChat. +- [Gradio Server](https://github.com/mit-han-lab/llm-awq/tree/nv_laptop/tinychat/serve#gradio-demo-vila-with-tinychat): Try to build your own VLM online demo with AWQ and TinyChat! +- [QServe](https://github.com/mit-han-lab/qserve): 🔥 **[New]** Efficient and accurate serving system for large-scale LLM inference. + +## Install + +1. Clone this repository and navigate to AWQ folder +``` +git clone https://github.com/mit-han-lab/llm-awq +cd llm-awq +``` + +2. Install Package +``` +conda create -n awq python=3.10 -y +conda activate awq +pip install --upgrade pip # enable PEP 660 support +pip install -e . +``` + +* For **edge devices** like Orin, before running the commands above, please: + + 1. Modify [pyproject.toml](pyproject.toml) by commenting out [this line](https://github.com/mit-han-lab/llm-awq/blob/3fce69061682fdd528824e5da3d03a8a8b545f2a/pyproject.toml#L17). + 2. Manually install precompiled PyTorch binaries (>=2.0.0) from [NVIDIA](https://forums.developer.nvidia.com/t/pytorch-for-jetson/72048). You also need to install torchvision from this website when running NVILA. + 3. Set the appropriate Python version for conda environment (e.g., `conda create -n awq python=3.8 -y` for JetPack 5). + +3. Install efficient W4A16 (4-bit weight, 16-bit activation) CUDA kernel and optimized FP16 kernels (e.g. layernorm, positional encodings). +``` +cd awq/kernels +python setup.py install +``` + +4. Install Flash Attention +``` +pip install flash-attn --no-build-isolation +``` + +We recommend starting an interactive python CLI interface and run `import flash_attn` to check whether FlashAttention-2 is installed successfully. If not, we recommend downloading pre-built wheels from [here](https://github.com/Dao-AILab/flash-attention/releases/tag/v2.5.8). Please notice: + +- PyTorch version needs to exactly match with the version specified in the `.whl` name; +- Check out both `cxx11abiTRUE` and `cxx11abiFALSE` wheels if one of them does not work; +- It's recommended to match CUDA version specified in the `.whl` filename, but minor mismatches (e.g. 12.1 vs 12.2, or even 11.8 vs 12.2) usually do not matter. + + +5. [Optional] In order to run AWQ and TinyChat with NVILA model family, please install VILA: + +```bash +git clone https://github.com/NVlabs/VILA.git +cd VILA +pip install -e . +``` + +## AWQ Model Zoo + +We provide pre-computed AWQ search results for multiple model families, including LLaMA, OPT, Vicuna, and LLaVA. To get the pre-computed AWQ search results, run: + +```bash +# git lfs install # install git lfs if not already +git clone https://huggingface.co/datasets/mit-han-lab/awq-model-zoo awq_cache +``` + +The detailed support list: + +| Models | Sizes | INT4-g128 | INT3-g128 | +| ------ | --------------------------- | --------- | --------- | +| [DeepSeek-R1-Distill](/scripts/DeepSeek_R1_Distill_example.sh) | 1.5B/7B/8B | ✅ | | +| [Qwen-2.5](/scripts/qwen_example.sh) | 7B/72B | ✅ | | +| [NVILA](/scripts/nvila_example.sh) | 3B/8B | ✅ | | +| [VILA-1.5](/scripts/vila15_example.sh) | 3B/8B/13B/40B | ✅ | ✅ | +| [Llama3](/scripts/llama_example.sh) | 8B/70B | ✅ | ✅ | +| [VILA](/scripts/vila_example.sh) | 7B/13B | ✅ | | +| [Llama2](/scripts/llama_example.sh) | 7B/13B/70B | ✅ | ✅ | +| [LLaMA](/scripts/llama2_example.sh) | 7B/13B/30B/65B | ✅ | ✅ | +| [OPT](/scripts/opt_example.sh) | 125m/1.3B/2.7B/6.7B/13B/30B | ✅ | ✅ | +| [CodeLlama](/scripts/codellama_example.sh) | 7B/13B/34B | ✅ | ✅ | +| [StarCoder](/scripts/starcoder_example.sh) | 15.5B | ✅ | ✅ | +| [Vicuna-v1.1](/scripts/vicuna_example.sh) | 7B/13B | ✅ | | +| [LLaVA-v0](/scripts/llava_example.sh) | 13B | ✅ | | + +Note: We only list models that we have prepare the [AWQ searching results](https://huggingface.co/datasets/mit-han-lab/awq-model-zoo/tree/main) in the table above. AWQ also supports models such as LLaVA-v1.5 7B, and you may need to run the [AWQ search](#usage) on your own to quantize these models. For our latest VLM NVILA, quantized weights are available [here](https://huggingface.co/Efficient-Large-Model/NVILA-AWQ). + +## Examples + +AWQ can be easily applied to various LMs thanks to its good generalization, including instruction-tuned models and multi-modal LMs. It provides an easy-to-use tool to reduce the serving cost of LLMs. + +Here we provide two examples of AWQ application: Vicuna-7B (chatbot) and LLaVA-13B (visual reasoning) under `./examples` directory. AWQ can easily reduce the GPU memory of model serving and speed up token generation. It provides accurate quantization, providing reasoning outputs. You should be able to observe **memory savings** when running the models with 4-bit weights. + +Note that we perform AWQ using only textual calibration data, depsite we are running on multi-modal input. Please refer to `./examples` for details. + +![overview](figures/example_vis.jpg) + +## Usage + +We provide several sample script to run AWQ (please refer to `./scripts`). We use Llama3-8B as an example. + +1. Perform AWQ search and save search results (we already did it for you): +```bash +python -m awq.entry --model_path /PATH/TO/LLAMA3/llama3-8b \ + --w_bit 4 --q_group_size 128 \ + --run_awq --dump_awq awq_cache/llama3-8b-w4-g128.pt +``` + +2. Evaluate the AWQ quantized model on WikiText-2 (simulated pseudo quantization) +```bash +python -m awq.entry --model_path /PATH/TO/LLAMA3/llama3-8b \ + --tasks wikitext \ + --w_bit 4 --q_group_size 128 \ + --load_awq awq_cache/llama3-8b-w4-g128.pt \ + --q_backend fake +``` + +3. Generate real quantized weights (INT4) +```bash +mkdir quant_cache +python -m awq.entry --model_path /PATH/TO/LLAMA3/llama3-8b \ + --w_bit 4 --q_group_size 128 \ + --load_awq awq_cache/llama3-8b-w4-g128.pt \ + --q_backend real --dump_quant quant_cache/llama3-8b-w4-g128-awq.pt +``` + +4. Load and evaluate the real quantized model (now you can see smaller gpu memory usage) +```bash +python -m awq.entry --model_path /PATH/TO/LLAMA3/llama3-8b \ + --tasks wikitext \ + --w_bit 4 --q_group_size 128 \ + --load_quant quant_cache/llama3-8b-w4-g128-awq.pt +``` +## Results on Visual Language Models + +AWQ also seamlessly supports large multi-modal models (LMMs). Please refer to [TinyChat](./tinychat/README.md) for more details. + + + + + + + +## Reference + +If you find AWQ useful or relevant to your research, please kindly cite our paper: + +``` +@inproceedings{lin2023awq, + title={AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration}, + author={Lin, Ji and Tang, Jiaming and Tang, Haotian and Yang, Shang and Chen, Wei-Ming and Wang, Wei-Chen and Xiao, Guangxuan and Dang, Xingyu and Gan, Chuang and Han, Song}, + booktitle={MLSys}, + year={2024} +} +``` + +## Related Projects + +[SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models](https://github.com/mit-han-lab/smoothquant) + +[GPTQ: Accurate Post-training Compression for Generative Pretrained Transformers](https://arxiv.org/abs/2210.17323) + +[Vicuna and FastChat](https://github.com/lm-sys/FastChat#readme) + +[LLaVA: Large Language and Vision Assistant](https://github.com/haotian-liu/LLaVA) + +[VILA: On Pre-training for Visual Language Models](https://github.com/Efficient-Large-Model/VILA) + diff --git a/llm-awq/awq.egg-info/PKG-INFO b/llm-awq/awq.egg-info/PKG-INFO new file mode 100644 index 0000000000000000000000000000000000000000..78436830d5fae25cebd2141251820d03b462a789 --- /dev/null +++ b/llm-awq/awq.egg-info/PKG-INFO @@ -0,0 +1,319 @@ +Metadata-Version: 2.4 +Name: awq +Version: 0.1.0 +Summary: An efficient and accurate low-bit weight quantization(INT3/4) method for LLMs. +Classifier: Programming Language :: Python :: 3 +Classifier: License :: OSI Approved :: Apache Software License +Requires-Python: >=3.8 +Description-Content-Type: text/markdown +License-File: LICENSE +Requires-Dist: accelerate==0.34.2 +Requires-Dist: sentencepiece +Requires-Dist: tokenizers>=0.12.1 +Requires-Dist: torch==2.3.0 +Requires-Dist: torchvision==0.18.0 +Requires-Dist: transformers==4.46.0 +Requires-Dist: lm_eval==0.3.0 +Requires-Dist: texttable +Requires-Dist: toml +Requires-Dist: attributedict +Requires-Dist: protobuf +Requires-Dist: gradio==3.35.2 +Requires-Dist: gradio_client==0.2.9 +Requires-Dist: fastapi +Requires-Dist: uvicorn +Requires-Dist: pydantic==1.10.19 +Dynamic: license-file + +# AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration +[[Paper](https://arxiv.org/abs/2306.00978)][[Website](https://hanlab.mit.edu/projects/awq)] + +**Efficient and accurate** low-bit weight quantization (INT3/4) for LLMs, supporting **instruction-tuned** models and **multi-modal** LMs. + +![overview](figures/overview.png) + +The current release supports: + +- AWQ search for accurate quantization. +- Pre-computed AWQ model zoo for LLMs (Llama-1/2/3, OPT, CodeLlama, StarCoder, Vicuna, VILA, LLaVA; load to generate quantized weights). +- Memory-efficient 4-bit Linear in PyTorch. +- Efficient CUDA kernel implementation for fast inference (support context and decoding stage). +- Examples on 4-bit inference of an instruction-tuned model (Vicuna) and **multi-modal LM** (VILA). +- Chunk prefilling for faster prefilling in multi-round Q&A setting. +- State-of-the-art prefilling speed of LLMs/VLMs on edge devices: [TinyChat 2.0](./tinychat). + +**Thanks to AWQ, TinyChat can deliver more efficient responses with LLM/VLM chatbots through 4-bit inference.** + +* TinyChat with LLaMA-3-8b on RTX 4090 (2.7x faster than FP16): + +![TinyChat with LLaMA-3-8b on RTX 4090: W4A16 is 2.7x faster than FP16](./tinychat/figures/4090_example_new.gif) + +* TinyChat with LLaMA-3-8b on Jetson Orin (2.9x faster than FP16): + +![TinyChat with LLaMA-3-8b on Jetson Orin: W4A16 is 2.9x faster than FP16](./tinychat/figures/orin_example_new.gif) + + +**TinyChat also supports inference with vision language models (e.g., VILA, LLaVA). In the following examples, W4A16 quantized models from VILA family are launched with TinyChat.** + +* TinyChat with NVILA-8B on RTX 4090 (single-image inputs): + +![TinyChat with NVILA on 4090 single image](./tinychat/figures/4090_nvila_single.gif) + +* TinyChat with NVILA-8B on RTX 4090 (multi-image inputs): + +![TinyChat with NVILA on 4090 multiple images](./tinychat/figures/4090_nvila_multi.gif) + + + +* TinyChat with video reasoning: + +https://github.com/user-attachments/assets/b68a7a0d-5175-4030-985b-5ae0ae94f874 + +**Prompt:** What might be the next step according to the video? + +**Answer:** The next step in the video could be to place the shaped dough onto a baking sheet and let it rise before baking. + +**Online demo:** https://vila.hanlab.ai + +Check out [TinyChat](tinychat), which offers a turn-key solution for **on-device inference** of LLMs and VLMs on **resource-constrained edge platforms**. With TinyChat, it is now possible to efficiently run **large** models on **small** and **low-power** devices even without Internet connection! + + +## News +- [2025/04] 🔥 AWQ now supports DeepSeek-R1-Distilled models. Try our example [here](https://github.com/mit-han-lab/llm-awq/blob/main/scripts/DeepSeek_R1_Distill_example.sh)! +- [2025/02] AWQ now supports BF16 precision. See example [here](https://github.com/mit-han-lab/llm-awq/blob/main/scripts/qwen_example.sh). +- [2024/10] 🔥⚡ Explore advancements in [TinyChat 2.0](./tinychat), the latest version with significant advancements in prefilling speed of Edge LLMs and VLMs, **1.5-1.7x** faster than the previous version of TinyChat. Please refer to the [README](./tinychat/README.md) and [blog](https://hanlab.mit.edu/blog/tinychat20) for more details. +- [2024/05] 🏆 AWQ receives the **Best Paper Award** at **MLSys 2024**. 🎉 +- [2024/05] 🔥 The **VILA-1.5** model family which features **video understanding** is now supported in AWQ and TinyChat. Check out out online demo powered by TinyChat [here](https://vila.hanlab.ai). Example is [here](scripts/vila15_example.sh). +- [2024/05] 🔥 [AMD](https://community.amd.com/t5/ai/reduce-memory-footprint-and-improve-performance-running-llms-on/ba-p/686157) adopts AWQ to improve LLM serving efficiency. +- [2024/04] 🔥 We released AWQ and TinyChat support for The **Llama-3** model family! Check out our example [here](scripts/llama3_example.sh). +- [2024/02] 🔥 AWQ has been accepted to **MLSys 2024**! +- [2024/02] 🔥 We supported [VILA Vision Languague Models](https://arxiv.org/abs/2312.07533) in AWQ & TinyChat! Check our latest demos with multi-image inputs! +- [2024/02] 🔥 We released new version of quantized GEMM/GEMV kernels in [**TinyChat**](tinychat), leading to **38 tokens/second** inference speed on NVIDIA Jetson Orin! +- [2024/01] 🔥 AWQ has been integrated by [Google Vertex AI](https://console.cloud.google.com/vertex-ai/publishers/meta/model-garden/llama-2-quantized)! +- [2023/11] 🔥 AWQ has been integrated by [Amazon Sagemaker Containers](https://aws.amazon.com/blogs/machine-learning/boost-inference-performance-for-llms-with-new-amazon-sagemaker-containers/)! +- [2023/11] 🔥 We added AWQ support and pre-computed search results for CodeLlama, StarCoder, StableCode models. Checkout our model zoo [here](https://huggingface.co/datasets/mit-han-lab/awq-model-zoo)! +- [2023/11] 🔥 AWQ is now integrated natively in Hugging Face transformers through `from_pretrained`. You can either load quantized models from the Hub or your own HF quantized models. +- [2023/10] AWQ is integrated into NVIDIA [TensorRT-LLM](https://github.com/NVIDIA/TensorRT-LLM/) +- [2023/09] AWQ is integrated into [Intel Neural Compressor](https://github.com/intel/neural-compressor), [FastChat](https://github.com/lm-sys/FastChat/blob/main/docs/awq.md), [vLLM](https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/quantization/awq.py), [HuggingFace TGI](https://github.com/huggingface/text-generation-inference/pull/1054), and [LMDeploy](https://github.com/InternLM/lmdeploy). +- [2023/09] ⚡ Check out our latest [**TinyChat**](tinychat), which is ~2x faster than the first release on Orin! +- [2023/09] ⚡ Check out [**AutoAWQ**](https://github.com/casper-hansen/AutoAWQ), a third-party implementation to make AWQ easier to expand to new models, improve inference speed, and integrate into Huggingface. +- [2023/07] 🔥 We released **TinyChat**, an efficient and lightweight chatbot interface based on AWQ. TinyChat enables efficient LLM inference on both cloud and edge GPUs. Llama-2-chat models are supported! Check out our implementation [here](tinychat). +- [2023/07] 🔥 We added AWQ support and pre-computed search results for Llama-2 models (7B & 13B). Checkout our model zoo [here](https://huggingface.co/datasets/mit-han-lab/awq-model-zoo)! +- [2023/07] We extended the support for more LLM models including MPT, Falcon, and BLOOM. + +## Contents + +- [AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration](#awq-activation-aware-weight-quantization-for-llm-compression-and-acceleration) + - [News](#news) + - [Contents](#contents) + - [Helpful Links](#helpful-links) + - [Install](#install) + - [AWQ Model Zoo](#awq-model-zoo) + - [Examples](#examples) + - [Usage](#usage) + - [Results on Visual Language Models](#results-on-visual-language-models) + - [Reference](#reference) + - [Related Projects](#related-projects) + +## Helpful Links + +- [VILA online demo](vila.hanlab.ai): Visual Language Models efficiently supported by AWQ & TinyChat. +- [LLM on the Edge](https://github.com/mit-han-lab/llm-awq/tree/nv_laptop?tab=readme-ov-file#install): AWQ and TinyChat support edge GPUs such as NVIDIA Jetson Orin. +- [VLMs on Laptop](https://github.com/mit-han-lab/llm-awq/tree/nv_laptop?tab=readme-ov-file#run-vila-on-laptop): Follow the instructions to deploy VLMs on NVIDIA Laptops with TinyChat. +- [Gradio Server](https://github.com/mit-han-lab/llm-awq/tree/nv_laptop/tinychat/serve#gradio-demo-vila-with-tinychat): Try to build your own VLM online demo with AWQ and TinyChat! +- [QServe](https://github.com/mit-han-lab/qserve): 🔥 **[New]** Efficient and accurate serving system for large-scale LLM inference. + +## Install + +1. Clone this repository and navigate to AWQ folder +``` +git clone https://github.com/mit-han-lab/llm-awq +cd llm-awq +``` + +2. Install Package +``` +conda create -n awq python=3.10 -y +conda activate awq +pip install --upgrade pip # enable PEP 660 support +pip install -e . +``` + +* For **edge devices** like Orin, before running the commands above, please: + + 1. Modify [pyproject.toml](pyproject.toml) by commenting out [this line](https://github.com/mit-han-lab/llm-awq/blob/3fce69061682fdd528824e5da3d03a8a8b545f2a/pyproject.toml#L17). + 2. Manually install precompiled PyTorch binaries (>=2.0.0) from [NVIDIA](https://forums.developer.nvidia.com/t/pytorch-for-jetson/72048). You also need to install torchvision from this website when running NVILA. + 3. Set the appropriate Python version for conda environment (e.g., `conda create -n awq python=3.8 -y` for JetPack 5). + +3. Install efficient W4A16 (4-bit weight, 16-bit activation) CUDA kernel and optimized FP16 kernels (e.g. layernorm, positional encodings). +``` +cd awq/kernels +python setup.py install +``` + +4. Install Flash Attention +``` +pip install flash-attn --no-build-isolation +``` + +We recommend starting an interactive python CLI interface and run `import flash_attn` to check whether FlashAttention-2 is installed successfully. If not, we recommend downloading pre-built wheels from [here](https://github.com/Dao-AILab/flash-attention/releases/tag/v2.5.8). Please notice: + +- PyTorch version needs to exactly match with the version specified in the `.whl` name; +- Check out both `cxx11abiTRUE` and `cxx11abiFALSE` wheels if one of them does not work; +- It's recommended to match CUDA version specified in the `.whl` filename, but minor mismatches (e.g. 12.1 vs 12.2, or even 11.8 vs 12.2) usually do not matter. + + +5. [Optional] In order to run AWQ and TinyChat with NVILA model family, please install VILA: + +```bash +git clone https://github.com/NVlabs/VILA.git +cd VILA +pip install -e . +``` + +## AWQ Model Zoo + +We provide pre-computed AWQ search results for multiple model families, including LLaMA, OPT, Vicuna, and LLaVA. To get the pre-computed AWQ search results, run: + +```bash +# git lfs install # install git lfs if not already +git clone https://huggingface.co/datasets/mit-han-lab/awq-model-zoo awq_cache +``` + +The detailed support list: + +| Models | Sizes | INT4-g128 | INT3-g128 | +| ------ | --------------------------- | --------- | --------- | +| [DeepSeek-R1-Distill](/scripts/DeepSeek_R1_Distill_example.sh) | 1.5B/7B/8B | ✅ | | +| [Qwen-2.5](/scripts/qwen_example.sh) | 7B/72B | ✅ | | +| [NVILA](/scripts/nvila_example.sh) | 3B/8B | ✅ | | +| [VILA-1.5](/scripts/vila15_example.sh) | 3B/8B/13B/40B | ✅ | ✅ | +| [Llama3](/scripts/llama_example.sh) | 8B/70B | ✅ | ✅ | +| [VILA](/scripts/vila_example.sh) | 7B/13B | ✅ | | +| [Llama2](/scripts/llama_example.sh) | 7B/13B/70B | ✅ | ✅ | +| [LLaMA](/scripts/llama2_example.sh) | 7B/13B/30B/65B | ✅ | ✅ | +| [OPT](/scripts/opt_example.sh) | 125m/1.3B/2.7B/6.7B/13B/30B | ✅ | ✅ | +| [CodeLlama](/scripts/codellama_example.sh) | 7B/13B/34B | ✅ | ✅ | +| [StarCoder](/scripts/starcoder_example.sh) | 15.5B | ✅ | ✅ | +| [Vicuna-v1.1](/scripts/vicuna_example.sh) | 7B/13B | ✅ | | +| [LLaVA-v0](/scripts/llava_example.sh) | 13B | ✅ | | + +Note: We only list models that we have prepare the [AWQ searching results](https://huggingface.co/datasets/mit-han-lab/awq-model-zoo/tree/main) in the table above. AWQ also supports models such as LLaVA-v1.5 7B, and you may need to run the [AWQ search](#usage) on your own to quantize these models. For our latest VLM NVILA, quantized weights are available [here](https://huggingface.co/Efficient-Large-Model/NVILA-AWQ). + +## Examples + +AWQ can be easily applied to various LMs thanks to its good generalization, including instruction-tuned models and multi-modal LMs. It provides an easy-to-use tool to reduce the serving cost of LLMs. + +Here we provide two examples of AWQ application: Vicuna-7B (chatbot) and LLaVA-13B (visual reasoning) under `./examples` directory. AWQ can easily reduce the GPU memory of model serving and speed up token generation. It provides accurate quantization, providing reasoning outputs. You should be able to observe **memory savings** when running the models with 4-bit weights. + +Note that we perform AWQ using only textual calibration data, depsite we are running on multi-modal input. Please refer to `./examples` for details. + +![overview](figures/example_vis.jpg) + +## Usage + +We provide several sample script to run AWQ (please refer to `./scripts`). We use Llama3-8B as an example. + +1. Perform AWQ search and save search results (we already did it for you): +```bash +python -m awq.entry --model_path /PATH/TO/LLAMA3/llama3-8b \ + --w_bit 4 --q_group_size 128 \ + --run_awq --dump_awq awq_cache/llama3-8b-w4-g128.pt +``` + +2. Evaluate the AWQ quantized model on WikiText-2 (simulated pseudo quantization) +```bash +python -m awq.entry --model_path /PATH/TO/LLAMA3/llama3-8b \ + --tasks wikitext \ + --w_bit 4 --q_group_size 128 \ + --load_awq awq_cache/llama3-8b-w4-g128.pt \ + --q_backend fake +``` + +3. Generate real quantized weights (INT4) +```bash +mkdir quant_cache +python -m awq.entry --model_path /PATH/TO/LLAMA3/llama3-8b \ + --w_bit 4 --q_group_size 128 \ + --load_awq awq_cache/llama3-8b-w4-g128.pt \ + --q_backend real --dump_quant quant_cache/llama3-8b-w4-g128-awq.pt +``` + +4. Load and evaluate the real quantized model (now you can see smaller gpu memory usage) +```bash +python -m awq.entry --model_path /PATH/TO/LLAMA3/llama3-8b \ + --tasks wikitext \ + --w_bit 4 --q_group_size 128 \ + --load_quant quant_cache/llama3-8b-w4-g128-awq.pt +``` +## Results on Visual Language Models + +AWQ also seamlessly supports large multi-modal models (LMMs). Please refer to [TinyChat](./tinychat/README.md) for more details. + + + + + + + +## Reference + +If you find AWQ useful or relevant to your research, please kindly cite our paper: + +``` +@inproceedings{lin2023awq, + title={AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration}, + author={Lin, Ji and Tang, Jiaming and Tang, Haotian and Yang, Shang and Chen, Wei-Ming and Wang, Wei-Chen and Xiao, Guangxuan and Dang, Xingyu and Gan, Chuang and Han, Song}, + booktitle={MLSys}, + year={2024} +} +``` + +## Related Projects + +[SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models](https://github.com/mit-han-lab/smoothquant) + +[GPTQ: Accurate Post-training Compression for Generative Pretrained Transformers](https://arxiv.org/abs/2210.17323) + +[Vicuna and FastChat](https://github.com/lm-sys/FastChat#readme) + +[LLaVA: Large Language and Vision Assistant](https://github.com/haotian-liu/LLaVA) + +[VILA: On Pre-training for Visual Language Models](https://github.com/Efficient-Large-Model/VILA) + diff --git a/llm-awq/awq/__pycache__/entry.cpython-311.pyc b/llm-awq/awq/__pycache__/entry.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..63d499804044a63a7b825a56224f5ba0d8f90dc9 Binary files /dev/null and b/llm-awq/awq/__pycache__/entry.cpython-311.pyc differ diff --git a/llm-awq/awq/entry.py b/llm-awq/awq/entry.py new file mode 100644 index 0000000000000000000000000000000000000000..134fc1c02a12ad5b34c3ab3e2e3ea528ccd5b1ae --- /dev/null +++ b/llm-awq/awq/entry.py @@ -0,0 +1,357 @@ +from lm_eval import evaluator, tasks +from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig +import torch +import argparse +import os +import json +from accelerate import ( + init_empty_weights, + infer_auto_device_map, + dispatch_model, + load_checkpoint_in_model, +) +from accelerate.utils.modeling import get_balanced_memory +from awq.utils.parallel import auto_parallel +from awq.quantize.pre_quant import run_awq, apply_awq +from awq.quantize.quantizer import ( + pseudo_quantize_model_weight, + real_quantize_model_weight, +) +from awq.utils.lm_eval_adaptor import LMEvalAdaptor +from awq.utils.utils import simple_dispatch_model +from datasets import load_dataset +from torch import nn +import tqdm + +parser = argparse.ArgumentParser() +parser.add_argument("--model_path", type=str, help="path of the hf model") +parser.add_argument("--dtype", type=str, default="float16", choices=["float16", "bfloat16"]) +parser.add_argument("--batch_size", type=int, default=1, help="batch size") +parser.add_argument("--tasks", default=None, type=str) +parser.add_argument("--output_path", default=None, type=str) +parser.add_argument("--num_fewshot", type=int, default=0) +# model config +parser.add_argument("--parallel", action="store_true", help="enable model parallelism") +# max memory to offload larger models to CPU +parser.add_argument( + "--max_memory", + type=str, + nargs="*", + help="List of device_id:max_memory pairs to be parsed into a dictionary; " + + "Example: 0:10GiB 1:10GiB cpu:30GiB; " + + "mode details here: " + + "https://huggingface.co/docs/accelerate/usage_guides/big_modeling", +) +parser.add_argument( + "--auto_parallel", + action="store_true", + help="automatically set parallel and batch_size", +) +# quantization config +parser.add_argument("--w_bit", type=int, default=None) +parser.add_argument("--q_group_size", type=int, default=-1) +parser.add_argument("--no_zero_point", action="store_true", help="disable zero_point") +parser.add_argument("--q_backend", type=str, default="fake", choices=["fake", "real"]) +# save/load real quantized weights +parser.add_argument("--dump_quant", type=str, default=None, help="save quantized model") +parser.add_argument( + "--dump_fake", type=str, default=None, help="save fake-quantized model" +) +parser.add_argument("--load_quant", type=str, default=None, help="load quantized model") +# apply/save/load awq +parser.add_argument("--run_awq", action="store_true", help="perform awq search process") +parser.add_argument( + "--dump_awq", type=str, default=None, help="save the awq search results" +) +parser.add_argument( + "--load_awq", type=str, default=None, help="load the awq search results" +) +parser.add_argument( + "--vila-15", + action="store_true", + help="quantizing vila 1.5", +) +parser.add_argument( + "--vila-20", + action="store_true", + help="quantizing or smoothing vila 2.0 (NVILA)", +) +parser.add_argument( + "--smooth_scale", + action="store_true", + help="generate the act scale of visiontower", +) +parser.add_argument( + "--media_path", + type=str, + nargs="+", + help="The input video to get act scale for visiontower", +) +parser.add_argument( + "--act_scale_path", + type=str, + default=None, + help="Path to save act scale", +) +args = parser.parse_args() +assert ( + args.act_scale_path is not None and len(args.media_path) > 0 +) or not args.smooth_scale +vila_10_quant_mode = ( + ("llava" in args.model_path.lower() or "vila" in args.model_path.lower()) + and not args.vila_15 + and not args.vila_20 +) + +max_memory = [v.split(":") for v in (args.max_memory or [])] +max_memory = {(int(k) if k.isdigit() else k): v for k, v in max_memory} + +if args.auto_parallel: + gpu_list = auto_parallel(args) + +# get quantization config (apart from w_bit) +q_config = { + "zero_point": not args.no_zero_point, # by default True + "q_group_size": args.q_group_size, # whether to use group quantization +} +print("Quantization config:", q_config) + +# build model and tokenizer + + +def build_model_and_enc(model_path, dtype): + torch_dtype = torch.float16 if dtype == "float16" else torch.bfloat16 + if not os.path.exists(model_path): # look into ssd + raise FileNotFoundError(f"{model_path} not found!") + print(f"* Building model {model_path}") + + # all hf model + if vila_10_quant_mode: + from llava.model.builder import load_pretrained_model + from llava.mm_utils import get_model_name_from_path + + enc, model, image_processor, context_len = load_pretrained_model( + model_path=model_path, + model_base=None, + model_name=get_model_name_from_path(model_path), + device="cpu", + **{"use_cache": False}, + ) + else: + config = AutoConfig.from_pretrained(model_path, trust_remote_code=True) + # Note (Haotian): To avoid OOM after huggingface transformers 4.36.2 + config.use_cache = False + if "mpt" in config.__class__.__name__.lower(): + enc = AutoTokenizer.from_pretrained( + config.tokenizer_name, trust_remote_code=True + ) + else: + enc = AutoTokenizer.from_pretrained( + model_path, use_fast=False, trust_remote_code=True + ) + + if args.load_quant: # directly load quantized weights + print("Loading pre-computed quantized weights...") + with init_empty_weights(): + model = AutoModelForCausalLM.from_config( + config=config, torch_dtype=torch_dtype, trust_remote_code=True + ) + real_quantize_model_weight( + model, w_bit=args.w_bit, q_config=q_config, init_only=True + ) + + model.tie_weights() + + # Infer device map + kwargs = {"max_memory": max_memory} if len(max_memory) else {} + device_map = infer_auto_device_map( + model, + no_split_module_classes=[ + "OPTDecoderLayer", + "LlamaDecoderLayer", + "BloomBlock", + "MPTBlock", + "DecoderLayer", + ], + **kwargs, + ) + # Load checkpoint in the model + load_checkpoint_in_model( + model, + checkpoint=args.load_quant, + device_map=device_map, + offload_state_dict=True, + ) + # Dispatch model + model = simple_dispatch_model(model, device_map=device_map) + + model.eval() + else: # fp16 to quantized + args.run_awq &= not args.load_awq # if load_awq, no need to run awq + # Init model on CPU: + kwargs = {"torch_dtype": torch_dtype, "low_cpu_mem_usage": True} + if not vila_10_quant_mode: + model = AutoModelForCausalLM.from_pretrained( + model_path, config=config, trust_remote_code=True, **kwargs + ) + + model.eval() + + if args.run_awq: + assert args.dump_awq, "Please save the awq results with --dump_awq" + + awq_results = run_awq( + model, + enc, + w_bit=args.w_bit, + q_config=q_config, + n_samples=128, + seqlen=512, + ) + if args.dump_awq: + dirpath = os.path.dirname(args.dump_awq) + os.makedirs(dirpath, exist_ok=True) + + torch.save(awq_results, args.dump_awq) + print("AWQ results saved at", args.dump_awq) + + exit(0) + + if args.load_awq: + print("Loading pre-computed AWQ results from", args.load_awq) + awq_results = torch.load(args.load_awq, map_location="cpu") + apply_awq(model, awq_results) + + # weight quantization + if args.w_bit is not None: + if args.q_backend == "fake": + assert ( + args.dump_quant is None + ), "Need to use real quantization to dump quantized weights" + pseudo_quantize_model_weight(model, w_bit=args.w_bit, q_config=q_config) + if args.dump_fake: + model.save_pretrained(args.dump_fake) + print("Pseudo-quantized models saved at", args.dump_fake) + elif args.q_backend == "real": # real quantization + real_quantize_model_weight(model, w_bit=args.w_bit, q_config=q_config) + if args.dump_quant: + if not args.dump_quant.endswith("v2.pt"): + print("[Info] Auto-change the dump_quant file name to *v2.pt") + args.dump_quant = args.dump_quant.replace(".pt", "-v2.pt") + dirpath = os.path.dirname(args.dump_quant) + os.makedirs(dirpath, exist_ok=True) + + print(f"Saving the quantized model at {args.dump_quant}...") + torch.save(model.cpu().state_dict(), args.dump_quant) + exit(0) + else: + raise NotImplementedError + + # Move the model to GPU (as much as possible) for LM evaluation + kwargs = { + "max_memory": get_balanced_memory( + model, max_memory if len(max_memory) > 0 else None + ) + } + device_map = infer_auto_device_map( + model, + # TODO: can we remove this? + no_split_module_classes=[ + "OPTDecoderLayer", + "LlamaDecoderLayer", + "BloomBlock", + "MPTBlock", + "DecoderLayer", + ], + **kwargs, + ) + model = dispatch_model(model, device_map=device_map) + + return model, enc + + +def main(): + if args.output_path is not None and os.path.exists(args.output_path): + # print(f"Results {args.output_path} already generated. Exit.") + print(f"Results {args.output_path} already generated. Overwrite.") + # exit() + + # a hack here to auto set model group + if args.smooth_scale and args.vila_20: + if os.path.exists(args.act_scale_path): + print(f"Found existing Smooth Scales {args.act_scale_path}, skip.") + else: + from awq.quantize import get_smooth_scale + + act_scale = get_smooth_scale(args.model_path, args.media_path) + os.makedirs(os.path.dirname(args.act_scale_path), exist_ok=True) + torch.save(act_scale, args.act_scale_path) + print("Save act scales at " + str(args.act_scale_path)) + args.model_path = args.model_path + "/llm" + if args.dump_awq is None and args.dump_quant is None: + exit() + + if args.dump_awq and os.path.exists(args.dump_awq): + print(f"Found existing AWQ results {args.dump_awq}, exit.") + exit() + model, enc = build_model_and_enc(args.model_path, args.dtype) + + if args.tasks is not None: + # https://github.com/IST-DASLab/gptq/blob/2d65066eeb06a5c9ff5184d8cebdf33662c67faf/llama.py#L206 + if args.tasks == "wikitext": + testenc = load_dataset("wikitext", "wikitext-2-raw-v1", split="test") + testenc = enc("\n\n".join(testenc["text"]), return_tensors="pt") + model.seqlen = 2048 + testenc = testenc.input_ids.to(model.device) + nsamples = testenc.numel() // model.seqlen + model = model.eval() + nlls = [] + for i in tqdm.tqdm(range(nsamples), desc="evaluating..."): + batch = testenc[:, (i * model.seqlen) : ((i + 1) * model.seqlen)].to( + model.device + ) + with torch.no_grad(): + lm_logits = model(batch).logits + shift_logits = lm_logits[:, :-1, :].contiguous().float() + shift_labels = testenc[ + :, (i * model.seqlen) : ((i + 1) * model.seqlen) + ][:, 1:] + loss_fct = nn.CrossEntropyLoss() + loss = loss_fct( + shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1) + ) + neg_log_likelihood = loss.float() * model.seqlen + nlls.append(neg_log_likelihood) + + ppl = torch.exp(torch.stack(nlls).sum() / (nsamples * model.seqlen)) + print(ppl.item()) + + results = {"ppl": ppl.item()} + if args.output_path is not None: + os.makedirs(os.path.dirname(args.output_path), exist_ok=True) + with open(args.output_path, "w") as f: + json.dump(results, f, indent=2) + else: + task_names = args.tasks.split(",") + + lm_eval_model = LMEvalAdaptor(args.model_path, model, enc, args.batch_size) + results = evaluator.simple_evaluate( + model=lm_eval_model, + tasks=task_names, + batch_size=args.batch_size, + no_cache=True, + num_fewshot=args.num_fewshot, + ) + + print(evaluator.make_table(results)) + + if args.output_path is not None: + os.makedirs(os.path.dirname(args.output_path), exist_ok=True) + # otherwise cannot save + results["config"]["model"] = args.model_path + with open(args.output_path, "w") as f: + json.dump(results, f, indent=2) + + +if __name__ == "__main__": + main() diff --git a/llm-awq/awq/kernels/csrc/attention/README.md b/llm-awq/awq/kernels/csrc/attention/README.md new file mode 100644 index 0000000000000000000000000000000000000000..ec0aae55e81684ebce5f0d5fd13680470ced7845 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/attention/README.md @@ -0,0 +1,8 @@ +# Attention kernel from FasterTransformer + +This CUDA extension wraps the single-query attention [kernel](https://github.com/NVIDIA/FasterTransformer/blob/release/v5.2.1_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention/decoder_masked_multihead_attention_template.hpp) from +FasterTransformer v5.2.1 for benchmarking purpose. + +```sh +cd csrc/ft_attention && pip install . +``` diff --git a/llm-awq/awq/kernels/csrc/attention/cuda_bf16_fallbacks.cuh b/llm-awq/awq/kernels/csrc/attention/cuda_bf16_fallbacks.cuh new file mode 100644 index 0000000000000000000000000000000000000000..f5641f61609172090da1c8e77e43f9f4694ccca0 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/attention/cuda_bf16_fallbacks.cuh @@ -0,0 +1,257 @@ +// Downloaded from from FasterTransformer v5.2.1 +// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.2.1_tag/src/fastertransformer/utils/cuda_bf16_fallbacks.cuh +/* + * Copyright (c) 2019-2022, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "cuda_bf16_wrapper.h" +#include + +namespace fastertransformer { + +#ifdef ENABLE_BF16 +inline __device__ float2 bf1622float2(const __nv_bfloat162 val) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + float2 f_val; + f_val.x = __low2float(val); + f_val.y = __high2float(val); + return f_val; +#else + return __bfloat1622float2(val); +#endif +} + +inline __device__ int16_t bf1622int16(__nv_bfloat162 val) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + float2 f_val; + f_val.x = max(min(__low2float(val), 127.f), -128.f); + f_val.y = max(min(__high2float(val), 127.f), -128.f); + union { int8_t int8[2]; int16_t int16; }; + int8[0] = static_cast(static_cast(f_val.x)); + int8[1] = static_cast(static_cast(f_val.y)); + return int16; +#else + val = __hmin2(val, make_bfloat162(127., 127.)); + val = __hmax2(val, make_bfloat162(-128., -128.)); + union { int8_t int8[2]; int16_t int16; }; + int8[0] = static_cast(static_cast(val.x)); + int8[1] = static_cast(static_cast(val.y)); + return int16; +#endif +} + +inline __device__ __nv_bfloat162 float22bf162(const float2 val) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + return __floats2bfloat162_rn(val.x, val.y); +#else + return __float22bfloat162_rn(val); +#endif +} + +inline __device__ __nv_bfloat162 bf162bf162(const __nv_bfloat16 val) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + __nv_bfloat162 val2; + val2.x = val; + val2.y = val; + return val2; +#else + return __bfloat162bfloat162(val); +#endif +} + +inline __device__ __nv_bfloat162 bf16hadd2(const __nv_bfloat162 x, const __nv_bfloat162 y) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + float fxl, fxh, fyl, fyh; + fxl = __low2float(x); + fxh = __high2float(x); + fyl = __low2float(y); + fyh = __high2float(y); + return __floats2bfloat162_rn(fxl + fyl, fxh + fyh); +#else + return __hadd2(x, y); +#endif +} + +inline __device__ __nv_bfloat16 bf16hadd(const __nv_bfloat16 x, const __nv_bfloat16 y) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + return __float2bfloat16( __bfloat162float(x) + __bfloat162float(y) ); +#else + return __hadd(x, y); +#endif +} + +inline __device__ __nv_bfloat162 bf16hsub2(const __nv_bfloat162 x, const __nv_bfloat162 y) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + float fxl, fxh, fyl, fyh; + fxl = __low2float(x); + fxh = __high2float(x); + fyl = __low2float(y); + fyh = __high2float(y); + return __floats2bfloat162_rn(fxl - fyl, fxh - fyh); +#else + return __hsub2(x, y); +#endif +} + +inline __device__ __nv_bfloat16 bf16hsub(const __nv_bfloat16 x, const __nv_bfloat16 y) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + return __float2bfloat16( __bfloat162float(x) - __bfloat162float(y) ); +#else + return __hsub(x, y); +#endif +} + +inline __device__ __nv_bfloat162 bf16hmul2(const __nv_bfloat162 x, const __nv_bfloat162 y) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + float fxl, fxh, fyl, fyh; + fxl = __low2float(x); + fxh = __high2float(x); + fyl = __low2float(y); + fyh = __high2float(y); + return __floats2bfloat162_rn(fxl * fyl, fxh * fyh); +#else + return __hmul2(x, y); +#endif +} + +inline __device__ __nv_bfloat16 bf16hmul(const __nv_bfloat16 x, const __nv_bfloat16 y) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + return __float2bfloat16( __bfloat162float(x) * __bfloat162float(y) ); +#else + return __hmul(x, y); +#endif +} + +inline __device__ __nv_bfloat162 bf16hfma2(const __nv_bfloat162 x, const __nv_bfloat162 y, const __nv_bfloat162 z) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + float fxl, fxh, fyl, fyh, fzl, fzh; + fxl = __low2float(x); + fxh = __high2float(x); + fyl = __low2float(y); + fyh = __high2float(y); + fzl = __low2float(z); + fzh = __high2float(z); + return __floats2bfloat162_rn(fxl * fyl + fzl, fxh * fyh + fzh); +#else + return __hfma2(x, y, z); +#endif +} + +inline __device__ __nv_bfloat16 bf16hfma(const __nv_bfloat16 x, const __nv_bfloat16 y, const __nv_bfloat16 z) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + return __float2bfloat16( __bfloat162float(x) * __bfloat162float(y) + __bfloat162float(z)); +#else + return __hfma(x, y, z); +#endif +} + +inline __device__ __nv_bfloat162 bf16exp2(const __nv_bfloat162 x) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + float fxl, fxh; + fxl = __low2float(x); + fxh = __high2float(x);; + return __floats2bfloat162_rn(expf(fxl), expf(fxh)); +#else + return h2exp(x); +#endif +} + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 800) +inline __device__ __nv_bfloat162 operator*(const __nv_bfloat162 x, const __nv_bfloat162 y) { return bf16hmul2(x, y); }; +inline __device__ __nv_bfloat162 operator+(const __nv_bfloat162 x, const __nv_bfloat162 y) { return bf16hadd2(x, y); }; + +inline __device__ __nv_bfloat162 make_bfloat162(const __nv_bfloat16 x, const __nv_bfloat16 y) +{ + __nv_bfloat162 t; t.x = x; t.y = y; return t; +} + +#endif + +inline __device__ __nv_bfloat16 bf16hadd(__nv_bfloat16 a, __nv_bfloat16 b, __nv_bfloat16 c) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + return __float2bfloat16(__bfloat162float(a) + __bfloat162float(b) + __bfloat162float(c)); +#else + return a + b + c; +#endif +} + +inline __device__ __nv_bfloat16 bf16hadd(__nv_bfloat16 a, __nv_bfloat16 b, __nv_bfloat16 c, __nv_bfloat16 d) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + return __float2bfloat16(__bfloat162float(a) + __bfloat162float(b) + __bfloat162float(c) + __bfloat162float(d)); +#else + return (__nv_bfloat16)((float)a + (float)b + (float)c + (float)d); +#endif +} + +inline __device__ __nv_bfloat162 bf16hadd2(__nv_bfloat162 a, __nv_bfloat162 b, __nv_bfloat162 c) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + float fal, fah, fbl, fbh, fcl, fch; + fal = __low2float(a); + fah = __high2float(a); + fbl = __low2float(b); + fbh = __high2float(b); + fcl = __low2float(c); + fch = __high2float(c); + return __floats2bfloat162_rn(fal + fbl + fcl, fah + fbh + fch); +#else + return a + b + c; +#endif +} + +inline __device__ __nv_bfloat16 bf16hmul(__nv_bfloat16 a, __nv_bfloat16 b, __nv_bfloat16 c) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + return __float2bfloat16(__bfloat162float(a) * __bfloat162float(b) * __bfloat162float(c)); +#else + return a * b * c; +#endif +} + +inline __device__ __nv_bfloat162 bf16hmul2(__nv_bfloat162 a, __nv_bfloat162 b, __nv_bfloat162 c) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + float fal, fah, fbl, fbh, fcl, fch; + fal = __low2float(a); + fah = __high2float(a); + fbl = __low2float(b); + fbh = __high2float(b); + fcl = __low2float(c); + fch = __high2float(c); + return __floats2bfloat162_rn(fal * fbl * fcl, fah * fbh * fch); +#else + return a * b * c; +#endif +} + +inline __device__ __nv_bfloat162 bf16hfma2(__nv_bfloat162 a, __nv_bfloat162 b, __nv_bfloat162 c, __nv_bfloat162 d) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + float fal, fah, fbl, fbh, fcl, fch, fdl, fdh; + fal = __low2float(a); + fah = __high2float(a); + fbl = __low2float(b); + fbh = __high2float(b); + fcl = __low2float(c); + fch = __high2float(c); + fdl = __low2float(d); + fdh = __high2float(d); + return __floats2bfloat162_rn(fal * fbl * fcl + fdl, fah * fbh * fch + fdh); +#else + return a * b * c + d; +#endif +} + +#endif // ENABLE_BF16 + +} // namespace fastertransformer diff --git a/llm-awq/awq/kernels/csrc/attention/cuda_bf16_wrapper.h b/llm-awq/awq/kernels/csrc/attention/cuda_bf16_wrapper.h new file mode 100644 index 0000000000000000000000000000000000000000..efb6e798730879bc2cd16088b2091991862a6074 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/attention/cuda_bf16_wrapper.h @@ -0,0 +1,23 @@ +// Downloaded from from FasterTransformer v5.2.1 +// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.2.1_tag/src/fastertransformer/utils/cuda_bf16_wrapper.h +/* + * Copyright (c) 2019-2022, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#ifdef ENABLE_BF16 +#include +#endif diff --git a/llm-awq/awq/kernels/csrc/attention/decoder_masked_multihead_attention.cu b/llm-awq/awq/kernels/csrc/attention/decoder_masked_multihead_attention.cu new file mode 100644 index 0000000000000000000000000000000000000000..e5a6690086489a46c39452d7c9d3d14c7edf2ddf --- /dev/null +++ b/llm-awq/awq/kernels/csrc/attention/decoder_masked_multihead_attention.cu @@ -0,0 +1,154 @@ +// Adapted from from FasterTransformer v5.2.1 +// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.2.1_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention/decoder_masked_multihead_attention_128.cu +/* + * Copyright (c) 2020-2022, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "decoder_masked_multihead_attention.h" +#include "decoder_masked_multihead_attention_utils.h" +#include "cuda_bf16_wrapper.h" +#include +#include +#include + +#include "decoder_masked_multihead_attention_template.hpp" + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +#define MMHA_LAUNCH_KERNEL(T, Dh, Dh_MAX, THDS_PER_KEY, THDS_PER_VALUE, THDS_PER_BLOCK, DO_CROSS_ATTENTION, stream) \ + size_t smem_sz = mmha::smem_size_in_bytes(params, THDS_PER_VALUE, THDS_PER_BLOCK); \ + auto kernel = mmha::masked_multihead_attention_kernel; \ + if (smem_sz >= 48 * 1024) { \ + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_sz); \ + } \ + dim3 grid(params.num_heads, params.batch_size); \ + kernel<<>>(params) + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// !!! Specialize the launcher for Cross attention +template +void mmha_launch_kernel(const KERNEL_PARAMS_TYPE& params, const cudaStream_t& stream) +{ + constexpr int THREADS_PER_VALUE = Dh_MAX * sizeof(T) / 16; + constexpr bool DO_CROSS_ATTENTION = std::is_same>::value; + int tlength = (DO_CROSS_ATTENTION) ? params.memory_max_len : params.timestep; + // printf("tlength, CROSS_ATTENTION = %d, %d\n", tlength, DO_CROSS_ATTENTION); + if (tlength < 32) { + MMHA_LAUNCH_KERNEL(T, Dh, Dh_MAX, 4, THREADS_PER_VALUE, 64, DO_CROSS_ATTENTION, stream); + } + else if (tlength < 2048) { + MMHA_LAUNCH_KERNEL(T, Dh, Dh_MAX, 2, THREADS_PER_VALUE, 128, DO_CROSS_ATTENTION, stream); + } + else { + MMHA_LAUNCH_KERNEL(T, Dh, Dh_MAX, 1, THREADS_PER_VALUE, 256, DO_CROSS_ATTENTION, stream); + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +#undef MMHA_LAUNCH_KERNEL + +template +void multihead_attention_(const KERNEL_PARAMS_TYPE& params, const cudaStream_t& stream) +{ + switch (params.hidden_size_per_head) { + case 32: + mmha_launch_kernel(params, stream); + break; + case 48: + mmha_launch_kernel(params, stream); + break; + case 64: + mmha_launch_kernel(params, stream); + break; + case 80: + mmha_launch_kernel(params, stream); + break; + case 96: + mmha_launch_kernel(params, stream); + break; + case 112: + mmha_launch_kernel(params, stream); + break; + case 128: + mmha_launch_kernel(params, stream); + break; + case 160: + mmha_launch_kernel(params, stream); + break; + case 192: + mmha_launch_kernel(params, stream); + break; + case 224: + mmha_launch_kernel(params, stream); + break; + case 256: + mmha_launch_kernel(params, stream); + break; + default: + assert(false); + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +void masked_multihead_attention(const Masked_multihead_attention_params& params, const cudaStream_t& stream) +{ + multihead_attention_>(params, stream); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +void masked_multihead_attention(const Masked_multihead_attention_params& params, const cudaStream_t& stream) +{ + multihead_attention_>(params, stream); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +#ifdef ENABLE_BF16 +void masked_multihead_attention(const Masked_multihead_attention_params<__nv_bfloat16>& params, + const cudaStream_t& stream) +{ + multihead_attention_<__nv_bfloat16, Masked_multihead_attention_params<__nv_bfloat16>>(params, stream); +} +#endif +//////////////////////////////////////////////////////////////////////////////////////////////////// + +void cross_multihead_attention(const Cross_multihead_attention_params& params, const cudaStream_t& stream) +{ + multihead_attention_>(params, stream); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +void cross_multihead_attention(const Cross_multihead_attention_params& params, const cudaStream_t& stream) +{ + multihead_attention_>(params, stream); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +#ifdef ENABLE_BF16 +void cross_multihead_attention(const Cross_multihead_attention_params<__nv_bfloat16>& params, + const cudaStream_t& stream) +{ + multihead_attention_<__nv_bfloat16, Cross_multihead_attention_params<__nv_bfloat16>>(params, stream); +} +#endif + +//////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/llm-awq/awq/kernels/csrc/attention/decoder_masked_multihead_attention.h b/llm-awq/awq/kernels/csrc/attention/decoder_masked_multihead_attention.h new file mode 100644 index 0000000000000000000000000000000000000000..f68d48574900cd67487174784bd1bc26b037ac8b --- /dev/null +++ b/llm-awq/awq/kernels/csrc/attention/decoder_masked_multihead_attention.h @@ -0,0 +1,185 @@ +// Downloaded from from FasterTransformer v5.2.1 +// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.2.1_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention.h +/* + * Copyright (c) 2020-2022, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "cuda_bf16_wrapper.h" +#include +#include +#include +#include +#include + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +#define CHECK_CUDA(call) \ + do { \ + cudaError_t status_ = call; \ + if (status_ != cudaSuccess) { \ + fprintf(stderr, "CUDA error (%s:%d): %s\n", __FILE__, __LINE__, cudaGetErrorString(status_)); \ + exit(1); \ + } \ + } while (0) + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// The structure of parameters for the masked multihead attention kernel. +// +// We use the following terminology to describe the different dimensions. +// +// B: Batch size (number of sequences), +// L: Sequence length, +// D: Hidden dimension, +// H: Number of heads, +// Dh: Hidden dimension per head - Dh = D / H. + +template +struct Multihead_attention_params_base { + + // The output buffer. Dimensions B x D. + T* out = nullptr; + + // The input Qs and the associated bias. Dimensions B x D and D, resp. + const T *q = nullptr, *q_bias = nullptr; + // The input Ks and the associated bias. Dimensions B x D and D, resp. + const T *k = nullptr, *k_bias = nullptr; + // The input Vs and the associated bias. Dimensions B x D and D, resp. + const T *v = nullptr, *v_bias = nullptr; + + // The cache for the Ks. The size must be at least B x L x D. + T* k_cache = nullptr; + // The cache for the Vs. The size must be at least B x L x D. + T* v_cache = nullptr; + // The indirections to use for cache when beam sampling. + const int* cache_indir = nullptr; + + // Stride to handle the case when KQV is a single buffer + int stride = 0; + + // The batch size. + int batch_size = 0; + // The beam width + int beam_width = 0; + // The sequence length. + int memory_max_len = 0; + // The number of heads (H). + int num_heads = 0; + // The number of heads for KV cache. + int num_kv_heads = 0; + // The hidden dimension per head (Dh). + int hidden_size_per_head = 0; + // The per-head latent space reserved for rotary embeddings. + int rotary_embedding_dim = 0; + bool neox_rotary_style = false; + float rotary_base = 0.0f; + float rotary_scale = 1.0f; + // The maximum length of input sentences. + int max_input_length = 0; + // The current timestep. TODO(bhsueh) Check that do we only this param in cross attention? + int timestep = 0; + // The current timestep of each sentences (support different timestep for different sentences) + + // The 1.f / sqrt(Dh). Computed on the host. + float inv_sqrt_dh = 0.0f; + + // Used when we have some input context like gpt + const int* total_padding_tokens = nullptr; + + const bool* masked_tokens = nullptr; + const int* prefix_prompt_lengths = nullptr; + int max_prefix_prompt_length = 0; + + const T* relative_attention_bias = nullptr; + int relative_attention_bias_stride = 0; + // The slope per head of linear position bias to attention score (H). + const float* linear_bias_slopes = nullptr; + + const T* ia3_key_weights = nullptr; + const T* ia3_value_weights = nullptr; + const int* ia3_tasks = nullptr; + + const float* qkv_scale_out = nullptr; + const float* attention_out_scale = nullptr; + int int8_mode = 0; +}; + +template +struct Multihead_attention_params: public Multihead_attention_params_base { + // output cross attentions + float* cross_attention_out = nullptr; + int max_decoder_seq_len = 0; + bool is_return_cross_attentions = false; + + // allows to exist attention eary + bool* finished = nullptr; + + // required in case of cross attention + // will need it here till if constexpr in c++17 + int* memory_length_per_sample = nullptr; + + // required in case of masked attention with different length + const int* length_per_sample = nullptr; +}; + +template +struct Multihead_attention_params: public Multihead_attention_params_base { + // output cross attentions + float* cross_attention_out = nullptr; + int max_decoder_seq_len = 0; + bool is_return_cross_attentions = false; + + // allows to exist attention eary + bool* finished = nullptr; + + // required in case of cross attention + int* memory_length_per_sample = nullptr; + + // required in case of masked attention with different length + const int* length_per_sample = nullptr; +}; + +template +using Masked_multihead_attention_params = Multihead_attention_params; + +template +using Cross_multihead_attention_params = Multihead_attention_params; + +template +struct outputCrossAttentionParam { + // max decoder output length + int max_decoder_seq_len = 0; + T* cross_attention_out = nullptr; + bool is_return_cross_attentions = false; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +void masked_multihead_attention(const Masked_multihead_attention_params& params, const cudaStream_t& stream); +void masked_multihead_attention(const Masked_multihead_attention_params& params, const cudaStream_t& stream); +#ifdef ENABLE_BF16 +void masked_multihead_attention(const Masked_multihead_attention_params<__nv_bfloat16>& params, + const cudaStream_t& stream); +#endif +void cross_multihead_attention(const Cross_multihead_attention_params& params, const cudaStream_t& stream); +void cross_multihead_attention(const Cross_multihead_attention_params& params, const cudaStream_t& stream); +#ifdef ENABLE_BF16 +void cross_multihead_attention(const Cross_multihead_attention_params<__nv_bfloat16>& params, + const cudaStream_t& stream); +#endif + +//////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/llm-awq/awq/kernels/csrc/attention/decoder_masked_multihead_attention_template.hpp b/llm-awq/awq/kernels/csrc/attention/decoder_masked_multihead_attention_template.hpp new file mode 100644 index 0000000000000000000000000000000000000000..c336d8017e026f66d77af8e40f2bb0494ea61214 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/attention/decoder_masked_multihead_attention_template.hpp @@ -0,0 +1,1608 @@ +// Downloaded from from FasterTransformer v5.2.1 +// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.2.1_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention/decoder_masked_multihead_attention_template.hpp +/* + * Copyright (c) 2020-2022, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include "decoder_masked_multihead_attention.h" +#include "decoder_masked_multihead_attention_utils.h" +#include "cuda_bf16_wrapper.h" +#include "cuda_bf16_fallbacks.cuh" +#include +#include +#include + +// #define MMHA_USE_HMMA_FOR_REDUCTION + +// Below are knobs to extend FP32 accumulation for higher FP16 accuracy + +// Does not seem to affect the accuracy that much +#define MMHA_USE_FP32_ACUM_FOR_FMA + +// Seems to slightly improve the accuracy +#define MMHA_USE_FP32_ACUM_FOR_OUT + +#if 0 && defined(MMHA_USE_FP32_ACUM_FOR_OUT) + // Does not seem to improve the accuracy + //#define MMHA_USE_FP32_ACUM_FOR_LOGITS +#endif + +namespace mmha { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// +// We use the following terminology to describe the different dimensions. +// +// B: Batch size (number of sequences), +// L: Sequence length, +// D: Hidden dimension, +// H: Number of heads, +// Dh: Hidden dimension per head - Dh = D / H. +// +// The different kernels assign a threadblock for B x H pair. The grid has size (1, B, H). We use +// 64, 128 and 256 threads per block. +// +// Each threadblock loads Dh values from Q and its associated bias. The kernels run a loop to +// compute Q * K^T where K is loaded from a cache buffer -- except for the current timestep. The +// cache buffer helps with memory accesses and contains keys with bias. +// +// The layout of the cache buffer for the keys is [B, H, Dh/x, L, x] where x == 8 for FP16 and +// x == 4 for FP32 where the fastest moving dimension (contiguous data) is the rightmost one. The +// values for x are chosen to create chunks of 16 bytes. +// +// The different kernels use 1, 2 or 4 threads per key (THREADS_PER_KEY). The size of the LDGs +// depends on the number of threads per key. Each thread sums Dh / THREADS_PER_KEY elements. At +// the end of each iteration of the Q * K^T loop, we perform a reduction between lanes using an +// HMMA instruction (Tensor Core). Each Q * K^T valuey is stored in shared memory in FP32. +// +// After that loop, a parallel softmax is computed across the different Q * K^T values stored in +// shared memory. +// +// The kernel ends with a loop over the values in V. We use THREADS_PER_VALUE to control how many +// timesteps are computed by loop iteration. As with the keys, the values are read from a cache +// except for the current timestep. The layout of the cache buffer for the values is much simpler +// as it is [B, H, L, Dh]. +// + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct Qk_vec_ { +}; + +template<> +struct Qk_vec_ { + using Type = float; +}; +template<> +struct Qk_vec_ { + using Type = float2; +}; +template<> +struct Qk_vec_ { + using Type = float4; +}; +template<> +struct Qk_vec_ { + using Type = float4; +}; +template<> +struct Qk_vec_ { + using Type = uint32_t; +}; +template<> +struct Qk_vec_ { + using Type = uint32_t; +}; +template<> +struct Qk_vec_ { + using Type = uint2; +}; +template<> +struct Qk_vec_ { + using Type = uint4; +}; +#ifdef ENABLE_BF16 +template<> +struct Qk_vec_<__nv_bfloat16, 32> { + using Type = __nv_bfloat162; +}; +template<> +struct Qk_vec_<__nv_bfloat16, 64> { + using Type = __nv_bfloat162; +}; +template<> +struct Qk_vec_<__nv_bfloat16, 128> { + using Type = bf16_4_t; +}; +template<> +struct Qk_vec_<__nv_bfloat16, 256> { + using Type = bf16_8_t; +}; +#endif // ENABLE_BF16 +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct K_vec_ { +}; + +template<> +struct K_vec_ { + using Type = float; +}; +template<> +struct K_vec_ { + using Type = float2; +}; +template<> +struct K_vec_ { + using Type = float4; +}; +template<> +struct K_vec_ { + using Type = uint32_t; +}; +template<> +struct K_vec_ { + using Type = uint2; +}; +template<> +struct K_vec_ { + using Type = uint4; +}; +#ifdef ENABLE_BF16 +template<> +struct K_vec_<__nv_bfloat16, 4> { + using Type = __nv_bfloat162; +}; +template<> +struct K_vec_<__nv_bfloat16, 2> { + using Type = bf16_4_t; +}; +template<> +struct K_vec_<__nv_bfloat16, 1> { + using Type = bf16_8_t; +}; +#endif // ENABLE_BF16 +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct V_vec_ { +}; + +template<> +struct V_vec_ { + using Type = float; +}; +template<> +struct V_vec_ { + using Type = float2; +}; +template<> +struct V_vec_ { + using Type = float4; +}; +template<> +struct V_vec_ { + using Type = uint32_t; +}; +template<> +struct V_vec_ { + using Type = uint2; +}; +template<> +struct V_vec_ { + using Type = uint4; +}; +#ifdef ENABLE_BF16 +template<> +struct V_vec_<__nv_bfloat16, 2> { + using Type = __nv_bfloat162; +}; +template<> +struct V_vec_<__nv_bfloat16, 4> { + using Type = bf16_4_t; +}; +template<> +struct V_vec_<__nv_bfloat16, 8> { + using Type = bf16_8_t; +}; +#endif // ENABLE_BF16 +//////////////////////////////////////////////////////////////////////////////////////////////////// + +#ifdef MMHA_USE_FP32_ACUM_FOR_FMA +template +struct Qk_vec_acum_fp32_ { +}; + +template<> +struct Qk_vec_acum_fp32_ { + using Type = float; +}; +template<> +struct Qk_vec_acum_fp32_ { + using Type = float2; +}; +template<> +struct Qk_vec_acum_fp32_ { + using Type = float4; +}; +// template<> struct Qk_vec_acum_fp32_ { using Type = float; }; +template<> +struct Qk_vec_acum_fp32_ { + using Type = float2; +}; +template<> +struct Qk_vec_acum_fp32_ { + using Type = Float4_; +}; +template<> +struct Qk_vec_acum_fp32_ { + using Type = Float8_; +}; +template<> +struct Qk_vec_acum_fp32_<__nv_bfloat16> { + using Type = float; +}; +template<> +struct Qk_vec_acum_fp32_<__nv_bfloat162> { + using Type = float2; +}; +template<> +struct Qk_vec_acum_fp32_ { + using Type = Float4_; +}; +template<> +struct Qk_vec_acum_fp32_ { + using Type = Float8_; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct K_vec_acum_fp32_ { +}; + +template<> +struct K_vec_acum_fp32_ { + using Type = float; +}; +template<> +struct K_vec_acum_fp32_ { + using Type = float2; +}; +template<> +struct K_vec_acum_fp32_ { + using Type = float4; +}; +template<> +struct K_vec_acum_fp32_ { + using Type = float2; +}; +template<> +struct K_vec_acum_fp32_ { + using Type = Float4_; +}; +template<> +struct K_vec_acum_fp32_ { + using Type = Float8_; +}; +template<> +struct K_vec_acum_fp32_<__nv_bfloat16> { + using Type = float; +}; +template<> +struct K_vec_acum_fp32_<__nv_bfloat162> { + using Type = float2; +}; +template<> +struct K_vec_acum_fp32_ { + using Type = Float4_; +}; +template<> +struct K_vec_acum_fp32_ { + using Type = Float8_; +}; +#endif + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +#ifdef MMHA_USE_FP32_ACUM_FOR_OUT +template +struct V_vec_acum_fp32_ { +}; + +template<> +struct V_vec_acum_fp32_ { + using Type = float; +}; +template<> +struct V_vec_acum_fp32_ { + using Type = float2; +}; +template<> +struct V_vec_acum_fp32_ { + using Type = float4; +}; +template<> +struct V_vec_acum_fp32_ { + using Type = float2; +}; +template<> +struct V_vec_acum_fp32_ { + using Type = Float4_; +}; +template<> +struct V_vec_acum_fp32_ { + using Type = Float8_; +}; +#ifdef ENABLE_BF16 +template<> +struct V_vec_acum_fp32_<__nv_bfloat162> { + using Type = float2; +}; +template<> +struct V_vec_acum_fp32_ { + using Type = Float4_; +}; +template<> +struct V_vec_acum_fp32_ { + using Type = Float8_; +}; +#endif // ENABLE_BF16 +#endif +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +inline __device__ float qk_dot_(const K_vec (&q)[N], const K_vec (&k)[N]) +{ +#ifdef MMHA_USE_FP32_ACUM_FOR_FMA + using K_vec_acum = typename K_vec_acum_fp32_::Type; +#else + using K_vec_acum = K_vec; +#endif + // Compute the parallel products for Q*K^T (treat vector lanes separately). + K_vec_acum qk_vec = mul(q[0], k[0]); +#pragma unroll + for (int ii = 1; ii < N; ++ii) { + qk_vec = fma(q[ii], k[ii], qk_vec); + } + + // Finalize the reduction across lanes. + float qk = sum(qk_vec); +#pragma unroll + for (int mask = THREADS_PER_KEY / 2; mask >= 1; mask /= 2) { + qk += __shfl_xor_sync(uint32_t(-1), qk, mask); + } + return qk; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct Qk_dot { + template + static inline __device__ float dot(const K_vec (&q)[N], const K_vec (&k)[N]) + { + return qk_dot_(q, k); + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float4 hmma_fp32(const uint2& a, uint32_t b) +{ + float4 c; + float zero = 0.f; + asm volatile("mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 \n" + " {%0, %1, %2, %3}, \n" + " {%4, %5}, \n" + " {%6}, \n" + " {%7, %7, %7, %7}; \n" + + : "=f"(c.x), "=f"(c.y), "=f"(c.z), "=f"(c.w) + : "r"(a.x) "r"(a.y), "r"(b), "f"(zero)); + return c; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +inline __device__ float qk_hmma_dot_(const uint32_t (&q)[N], const uint32_t (&k)[N]) +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 750 +#ifdef MMHA_USE_FP32_ACUM_FOR_FMA + using K_vec_acum = typename K_vec_acum_fp32_::Type; +#else + using K_vec_acum = uint32_t; +#endif + K_vec_acum qk_vec = mul(q[0], k[0]); +#pragma unroll + for (int ii = 1; ii < N; ++ii) { + qk_vec = fma(q[ii], k[ii], qk_vec); + } +#ifdef MMHA_USE_FP32_ACUM_FOR_FMA + uint32_t qk_vec_ = float2_to_half2(qk_vec); + return hmma_fp32(make_uint2(qk_vec_, 0u), 0x3c003c00u).x; +#else + return hmma_fp32(make_uint2(qk_vec, 0u), 0x3c003c00u).x; +#endif +#else + return 0.f; +#endif +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template<> +struct Qk_dot { + template + static inline __device__ float dot(const uint32_t (&q)[N], const uint32_t (&k)[N]) + { +#if __CUDA_ARCH__ >= 750 && defined(MMHA_USE_HMMA_FOR_REDUCTION) + return qk_hmma_dot_(q, k); +#else + return qk_dot_<4>(q, k); +#endif // defined MMHA_USE_HMMA_FOR_REDUCTION + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +inline __device__ float block_sum(float* red_smem, float sum) +{ + + // Decompose the thread index into warp / lane. + int warp = threadIdx.x / WARP_SIZE; + int lane = threadIdx.x % WARP_SIZE; + +// Compute the sum per warp. +#pragma unroll + for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) { + sum += __shfl_xor_sync(uint32_t(-1), sum, mask); + } + + // Warp leaders store the data to shared memory. + if (lane == 0) { + red_smem[warp] = sum; + } + + // Make sure the data is in shared memory. + __syncthreads(); + + // The warps compute the final sums. + if (lane < WARPS_PER_BLOCK) { + sum = red_smem[lane]; + } + +// Parallel reduction inside the warp. +#pragma unroll + for (int mask = WARPS_PER_BLOCK / 2; mask >= 1; mask /= 2) { + sum += __shfl_xor_sync(uint32_t(-1), sum, mask); + } + + // Broadcast to other threads. + return __shfl_sync(uint32_t(-1), sum, 0); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ void convert_from_float(float& dst, float src) +{ + dst = src; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ void convert_from_float(uint16_t& dst, float src) +{ + dst = float_to_half(src); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ void convert_from_float(uint32_t& dst, float2 src) +{ + dst = float2_to_half2(src); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// +#ifdef ENABLE_BF16 +inline __device__ void convert_from_float(__nv_bfloat16& dst, float src) +{ + dst = __float2bfloat16(src); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ void convert_from_float(__nv_bfloat162& dst, float2 src) +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + dst = __float22bfloat162_rn(src); +#else + dst = __floats2bfloat162_rn(src.x, src.y); +#endif +} +#endif // ENABLE_BF16 +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ void convert_from_float(uint2& dst, Float4_ src) +{ + dst.x = float2_to_half2(src.x); + dst.y = float2_to_half2(src.y); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ void convert_from_float(uint2& dst, float4 src) +{ + convert_from_float(dst, Float4_{make_float2(src.x, src.y), make_float2(src.z, src.w)}); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ void convert_from_float(uint4& dst, Float8_ src) +{ + dst.x = float2_to_half2(src.x); + dst.y = float2_to_half2(src.y); + dst.z = float2_to_half2(src.z); + dst.w = float2_to_half2(src.w); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +#ifdef ENABLE_BF16 +inline __device__ void convert_from_float(bf16_4_t& dst, Float4_ src) +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + dst.x = __float22bfloat162_rn(src.x); + dst.y = __float22bfloat162_rn(src.y); +#else + dst.x = __floats2bfloat162_rn(src.x.x, src.x.y); + dst.y = __floats2bfloat162_rn(src.y.x, src.y.y); +#endif +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ void convert_from_float(bf16_4_t& dst, float4 src) +{ + convert_from_float(dst, Float4_{make_float2(src.x, src.y), make_float2(src.z, src.w)}); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ void convert_from_float(bf16_8_t& dst, Float8_ src) +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + dst.x = __float22bfloat162_rn(src.x); + dst.y = __float22bfloat162_rn(src.y); + dst.z = __float22bfloat162_rn(src.z); + dst.w = __float22bfloat162_rn(src.w); +#else + dst.x = __floats2bfloat162_rn(src.x.x, src.x.y); + dst.y = __floats2bfloat162_rn(src.y.x, src.y.y); + dst.z = __floats2bfloat162_rn(src.z.x, src.z.y); + dst.w = __floats2bfloat162_rn(src.w.x, src.w.y); +#endif +} +#endif // ENABLE_BF16 + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ void convert_from_float(float2& dst, float2 src) +{ + dst = src; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ void convert_from_float(float4& dst, float4 src) +{ + dst = src; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float convert_to_float(float4 u) +{ + return u.x; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float convert_to_float(uint4 u) +{ + float2 tmp = half2_to_float2(u.x); + return tmp.x; +} + +#if defined(MMHA_USE_FP32_ACUM_FOR_LOGITS) + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float cast_to_float(float u) +{ + return u; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float2 cast_to_float(float2 u) +{ + return u; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float4 cast_to_float(float4 u) +{ + return u; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ Float4_ cast_to_float(Float4_ u) +{ + return u; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ Float8_ cast_to_float(Float8_ u) +{ + return u; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float2 cast_to_float(uint32_t u) +{ + return half2_to_float2(u); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ Float4_ cast_to_float(uint2 u) +{ + Float4_ tmp; + tmp.x = half2_to_float2(u.x); + tmp.y = half2_to_float2(u.y); + return tmp; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ Float8_ cast_to_float(uint4 u) +{ + Float8_ tmp; + tmp.x = half2_to_float2(u.x); + tmp.y = half2_to_float2(u.y); + tmp.z = half2_to_float2(u.z); + tmp.w = half2_to_float2(u.w); + return tmp; +} + +#endif + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float float_from_int8(int8_t u) +{ + return u; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float2 float_from_int8(int16_t u) +{ + union { + int16_t int16; + int8_t int8[2]; + }; + int16 = u; + return make_float2(int8[0], int8[1]); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float4 float_from_int8(int32_t u) +{ + union { + int32_t int32; + int8_t int8[4]; + }; + int32 = u; + return make_float4(int8[0], int8[1], int8[2], int8[3]); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// clang-format off +inline __device__ Float8_ float_from_int8(int64_t u) +{ + union { + int64_t int64; + int16_t int16[4]; + }; + int64 = u; + return Float8_ {float_from_int8(int16[0]), + float_from_int8(int16[1]), + float_from_int8(int16[2]), + float_from_int8(int16[3])}; +} +// clang-format on + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ int8_t cast_to_int8(float val) +{ + union { + int8_t int8[2]; + int16_t int16; + }; + asm volatile("cvt.rni.sat.s8.f32 %0, %1;" : "=h"(int16) : "f"(val)); + return int8[0]; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ int32_t cast_to_int8(float4 val) +{ + union { + int8_t int8[4]; + int32_t int32; + }; + int8[0] = cast_to_int8(val.x); + int8[1] = cast_to_int8(val.y); + int8[2] = cast_to_int8(val.z); + int8[3] = cast_to_int8(val.w); + return int32; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ int64_t cast_to_int8(Float8_ val) +{ + union { + int8_t int8[8]; + int64_t int64; + }; + int8[0] = cast_to_int8(val.x.x); + int8[1] = cast_to_int8(val.x.y); + int8[2] = cast_to_int8(val.y.x); + int8[3] = cast_to_int8(val.y.y); + int8[4] = cast_to_int8(val.z.x); + int8[5] = cast_to_int8(val.z.y); + int8[6] = cast_to_int8(val.w.x); + int8[7] = cast_to_int8(val.w.y); + return int64; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +inline __device__ __host__ T div_up(T m, T n) +{ + return (m + n - 1) / n; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +inline size_t smem_size_in_bytes(const Multihead_attention_params& params, + int threads_per_value, + int threads_per_block) +{ + // The amount of shared memory needed to store the Q*K^T values in float. + const int max_timesteps = min(params.timestep, params.memory_max_len); + size_t qk_sz = (DO_CROSS_ATTENTION) ? div_up(params.memory_max_len + 1, 4) * 16 : div_up(max_timesteps + 1, 4) * 16; + + // The extra memory needed if we are not using floats for the final logits. + size_t logits_sz = 0; +#ifndef MMHA_USE_FP32_ACUM_FOR_LOGITS + if (sizeof(T) != 4) { + // TDOD + logits_sz = (DO_CROSS_ATTENTION) ? div_up(params.memory_max_len + 1, 4) * 4 * sizeof(T) : + div_up(max_timesteps + 1, 4) * 4 * sizeof(T); + } +#endif + + // The total size needed during softmax. + size_t softmax_sz = qk_sz + logits_sz; + + // The number of partial rows to reduce in the final reduction. + int rows_per_red = threads_per_block / threads_per_value; + // The amount of storage needed to finalize the outputs. + size_t red_sz = rows_per_red * params.hidden_size_per_head * sizeof(T) / 2; + + size_t transpose_rotary_size = 0; + if (params.rotary_embedding_dim > 0 && params.neox_rotary_style) { + transpose_rotary_size = 2 * params.rotary_embedding_dim * sizeof(T); + } + + // The max. + return max(max(softmax_sz, red_sz), transpose_rotary_size); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ constexpr uint32_t shfl_mask(int threads) +{ + return threads == 32 ? uint32_t(-1) : (1u << threads) - 1u; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template< + // The type of the inputs. Supported types: float and half. + typename T, + // The hidden dimension per head. + int Dh, + int Dh_MAX, + // The number of threads per key. + int THREADS_PER_KEY, + // The number of threads per value. + int THREADS_PER_VALUE, + // The number of threads in a threadblock. + int THREADS_PER_BLOCK, + bool DO_CROSS_ATTENTION> +__global__ void masked_multihead_attention_kernel(Multihead_attention_params params) +{ + + // Make sure the hidden dimension per head is a multiple of the number of threads per key. + static_assert(Dh_MAX % THREADS_PER_KEY == 0, ""); + // Make sure the hidden dimension per head is a multiple of the number of threads per value. + static_assert(Dh_MAX % THREADS_PER_VALUE == 0, ""); + + // The size of a warp. + constexpr int WARP_SIZE = 32; + // The number of warps in a threadblock. + constexpr int WARPS_PER_BLOCK = THREADS_PER_BLOCK / WARP_SIZE; + + // Use smem_size_in_bytes (above) to determine the amount of shared memory. + extern __shared__ char smem_[]; + + // The shared memory for the Q*K^T values and partial logits in softmax. + float* qk_smem = reinterpret_cast(smem_); + + // The shared memory for the logits. For FP32, that's the same buffer as qk_smem. + char* logits_smem_ = smem_; +#ifndef MMHA_USE_FP32_ACUM_FOR_LOGITS + if (sizeof(T) != 4) { + // TODO - change to tlength + const int max_timesteps = min(params.timestep, params.memory_max_len); + logits_smem_ += + (DO_CROSS_ATTENTION) ? div_up(params.memory_max_len + 1, 4) * 16 : div_up(max_timesteps + 1, 4) * 16; + } + T* logits_smem = reinterpret_cast(logits_smem_); +#else + float* logits_smem = reinterpret_cast(logits_smem_); +#endif + + // The shared memory to do the final reduction for the output values. Reuse qk_smem. + T* out_smem = reinterpret_cast(smem_); + + // The shared memory buffers for the block-wide reductions. One for max, one for sum. + __shared__ float red_smem[WARPS_PER_BLOCK * 2]; + + // A vector of Q or K elements for the current timestep. + using Qk_vec = typename Qk_vec_::Type; + + // Use alignment for safely casting the shared buffers as Qk_vec. + // Shared memory to store Q inputs. + __shared__ __align__(sizeof(Qk_vec)) T q_smem[Dh_MAX]; + + // This is one of the reasons we should have a separate kernel for cross attention + __shared__ __align__(sizeof(Qk_vec)) T bias_smem[DO_CROSS_ATTENTION ? Dh_MAX : 1]; + + // A vector of Q or K elements for the current timestep. + using Qk_vec = typename Qk_vec_::Type; + // The number of elements per vector. + constexpr int QK_VEC_SIZE = sizeof(Qk_vec) / sizeof(T); + // Make sure the hidden size per head is a multiple of the vector size. + static_assert(Dh_MAX % QK_VEC_SIZE == 0, ""); + // We will use block wide reduction if needed + // static_assert(Dh_MAX / QK_VEC_SIZE <= WARP_SIZE, ""); + // The number of vectors per warp. + constexpr int QK_VECS_PER_WARP = Dh_MAX / QK_VEC_SIZE; + + // The layout of the cache is [B, H, Dh/x, L, x] with x == 4/8 for FP32/FP16. Since each thread + // owns x elements, we have to decompose the linear index into chunks of x values and the posi- + // tion of the thread in that chunk. + + // The number of elements in a chunk of 16B (that's the x in the above formula). + constexpr int QK_ELTS_IN_16B = 16 / sizeof(T); + // The number of K vectors in 16B. + constexpr int QK_VECS_IN_16B = 16 / sizeof(Qk_vec); + + // The batch/beam idx + const int bi = blockIdx.y; + if (params.finished != nullptr && params.finished[bi] == true) { + return; + } + // The beam idx + const int beami = bi % params.beam_width; + // The "beam-aware" batch idx + const int bbi = bi / params.beam_width; + // The head. + const int num_kv_heads = params.num_kv_heads; + const int kv_rep = (params.num_heads / num_kv_heads); + const int hi = blockIdx.x; + const int hi_kv = hi / kv_rep; + + // Combine the batch and the head indices. + const int bhi = bi * params.num_heads + hi; + const int bhi_kv = bi * (params.num_heads / kv_rep) + hi_kv; + // Combine the "beam-aware" batch idx and the head indices. + const int bbhi = bbi * params.beam_width * params.num_heads + hi; + const int bbhi_kv = bbi * params.beam_width * (params.num_heads / kv_rep) + hi_kv; + // The thread in the block. + const int tidx = threadIdx.x; + + const bool handle_kv = !DO_CROSS_ATTENTION || (DO_CROSS_ATTENTION && params.timestep == 0); + // Every kv_rep threads have the same kv_cache values. So only the first one writes back. + const int write_kv_cache = handle_kv && (hi % kv_rep == 0); + + // While doing the product Q*K^T for the different keys we track the max. + float qk_max = -FLT_MAX; + + float qk = 0.0F; + + // int qkv_base_offset = (params.stride == 0) ? bhi * Dh : bi * params.stride + hi * Dh; + const int q_base_offset = bi * params.stride + hi * Dh; + const int k_base_offset = bi * params.stride + hi_kv * Dh; + const int v_base_offset = k_base_offset; + + const size_t bi_seq_len_offset = bi * params.memory_max_len; + + // int tlength = (DO_CROSS_ATTENTION)? params.memory_length_per_sample[bi] - 1 : params.timestep; + int tlength = (DO_CROSS_ATTENTION) ? params.memory_length_per_sample[bi] - 1 : + (params.length_per_sample == nullptr) ? + params.timestep : + params.length_per_sample[bi] + params.max_prefix_prompt_length; + const int first_step = max(0, tlength + 1 - params.memory_max_len); + const int tlength_circ = tlength % params.memory_max_len; + + // First QK_VECS_PER_WARP load Q and K + the bias values for the current timestep. + const bool is_masked = tidx >= QK_VECS_PER_WARP; + + // The offset in the Q and K buffer also accounts for the batch. + // int qk_offset = qkv_base_offset + tidx * QK_VEC_SIZE; + int q_offset = q_base_offset + tidx * QK_VEC_SIZE; + int k_offset = k_base_offset + tidx * QK_VEC_SIZE; + int v_offset = k_offset; + + // The offset in the bias buffer. + // int qk_bias_offset = hi * Dh + tidx * QK_VEC_SIZE; + int q_bias_offset = hi * Dh + tidx * QK_VEC_SIZE; + int k_bias_offset = hi_kv * Dh + tidx * QK_VEC_SIZE; + int v_bias_offset = k_bias_offset; + + const bool do_ia3 = handle_kv && params.ia3_tasks != nullptr; + const int ia3_task_id = do_ia3 ? params.ia3_tasks[bbi] : 0; + + // Trigger the loads from the Q and K buffers. + Qk_vec q; + zero(q); + if (!is_masked && (Dh == Dh_MAX || tidx * QK_VEC_SIZE < Dh)) { + if (params.int8_mode == 2) { + using Packed_Int8_t = typename packed_type::value>::type; + using Packed_Float_t = typename packed_type::value>::type; + const auto q_scaling = params.qkv_scale_out[0]; + const auto q_quant = + *reinterpret_cast(&reinterpret_cast(params.q)[q_offset]); + + convert_from_float(q, mul(q_scaling, float_from_int8(q_quant))); + } + else { + q = *reinterpret_cast(¶ms.q[q_offset]); + } + } + + Qk_vec k; + zero(k); + if (DO_CROSS_ATTENTION) { + // The 16B chunk written by the thread. + int co = tidx / QK_VECS_IN_16B; + // The position of the thread in that 16B chunk. + int ci = tidx % QK_VECS_IN_16B * QK_VEC_SIZE; + + // Two chunks are separated by L * x elements. A thread write QK_VEC_SIZE elements. + int offset = bhi_kv * params.memory_max_len * Dh + co * params.memory_max_len * QK_ELTS_IN_16B + + // params.timestep*QK_ELTS_IN_16B + + tlength * QK_ELTS_IN_16B + ci; + k = !is_masked && (Dh == Dh_MAX || tidx * QK_VEC_SIZE < Dh) ? + *reinterpret_cast(¶ms.k_cache[offset]) : + k; + } + else { + if (!is_masked && (Dh == Dh_MAX || tidx * QK_VEC_SIZE < Dh)) { + if (params.int8_mode == 2) { + using Packed_Int8_t = typename packed_type::value>::type; + using Packed_Float_t = typename packed_type::value>::type; + const auto k_scaling = params.qkv_scale_out[1]; + const auto k_quant = + *reinterpret_cast(&reinterpret_cast(params.k)[k_offset]); + + convert_from_float(k, mul(k_scaling, float_from_int8(k_quant))); + } + else { + k = *reinterpret_cast(¶ms.k[k_offset]); + } + } + } + + // Trigger the loads from the Q and K bias buffers. + Qk_vec q_bias; + zero(q_bias); + q_bias = (!is_masked && Dh == Dh_MAX || tidx * QK_VEC_SIZE < Dh) && params.q_bias != nullptr ? + *reinterpret_cast(¶ms.q_bias[q_bias_offset]) : + q_bias; + + Qk_vec k_bias; + zero(k_bias); + if (handle_kv) { + k_bias = !is_masked && (Dh == Dh_MAX || tidx * QK_VEC_SIZE < Dh) && params.k_bias != nullptr ? + *reinterpret_cast(¶ms.k_bias[k_bias_offset]) : + k_bias; + } + + // Computes the Q/K values with bias. + q = add(q, q_bias); + if (handle_kv) { + k = add(k, k_bias); + } + if (do_ia3 && !is_masked) { + k = mul( + k, + *reinterpret_cast( + ¶ms.ia3_key_weights[(ia3_task_id * params.num_heads + hi) * Dh + tidx * QK_VEC_SIZE])); + } + + // Padded len + const int padd_len = (params.total_padding_tokens == nullptr) ? 0 : params.total_padding_tokens[bi]; + if (params.rotary_embedding_dim > 0 && !params.neox_rotary_style) { + if (handle_kv) { + apply_rotary_embedding(q, k, tidx, params.rotary_embedding_dim, tlength - padd_len, params.rotary_base, params.rotary_scale); + } + else { + apply_rotary_embedding(q, tidx, params.rotary_embedding_dim, tlength - padd_len, params.rotary_base, params.rotary_scale); + } + } + else if (params.rotary_embedding_dim > 0 && params.neox_rotary_style) { + const bool do_rotary = !is_masked && QK_VEC_SIZE * tidx < params.rotary_embedding_dim; + + T* q_smem = reinterpret_cast(smem_); + T* k_smem = q_smem + params.rotary_embedding_dim; + + const int half_rotary_dim = params.rotary_embedding_dim / 2; + const int half_idx = (tidx * QK_VEC_SIZE) / half_rotary_dim; + const int intra_half_idx = (tidx * QK_VEC_SIZE) % half_rotary_dim; + const int smem_pitch = half_rotary_dim; // TODO: adjust for bank conflicts + + assert(half_rotary_dim % QK_VEC_SIZE == 0); + + if (do_rotary) { + *reinterpret_cast(q_smem + half_idx * smem_pitch + intra_half_idx) = q; + + if (handle_kv) { + *reinterpret_cast(k_smem + half_idx * smem_pitch + intra_half_idx) = k; + } + } + + __syncthreads(); + + const int transpose_idx = half_idx * (half_rotary_dim / 2) + intra_half_idx / 2; + constexpr int tidx_factor = (QK_VEC_SIZE > 1) ? QK_VEC_SIZE / 2 : 1; + if (do_rotary) { + mmha::vec_from_smem_transpose(q, q_smem, transpose_idx, smem_pitch); + + if (handle_kv) { + mmha::vec_from_smem_transpose(k, k_smem, transpose_idx, smem_pitch); + + mmha::apply_rotary_embedding( + q, k, transpose_idx / tidx_factor, params.rotary_embedding_dim, tlength - padd_len, params.rotary_base, params.rotary_scale); + + mmha::write_smem_transpose(k, k_smem, transpose_idx, smem_pitch); + } + else { + mmha::apply_rotary_embedding( + q, transpose_idx / tidx_factor, params.rotary_embedding_dim, tlength, params.rotary_base, params.rotary_scale); + } + mmha::write_smem_transpose(q, q_smem, transpose_idx, smem_pitch); + } + + __syncthreads(); + + if (do_rotary) { + q = *reinterpret_cast(q_smem + half_idx * smem_pitch + intra_half_idx); + if (handle_kv) { + k = *reinterpret_cast(k_smem + half_idx * smem_pitch + intra_half_idx); + } + } + + __syncthreads(); + } + + if (!is_masked) { + // Store the Q values to shared memory. + *reinterpret_cast(&q_smem[tidx * QK_VEC_SIZE]) = q; + + // Store Dh values of k_bias into smem, since will need to add later + // if params.timestep == 0 + if (DO_CROSS_ATTENTION && params.timestep == 0) { + *reinterpret_cast(&bias_smem[tidx * QK_VEC_SIZE]) = k_bias; + } + + // Write the K values to the global memory cache. + // + // NOTE: The stores are uncoalesced as we have multiple chunks of 16B spread across the memory + // system. We designed it this way as it allows much better memory loads (and there are many + // more loads) + the stores are really "write and forget" since we won't need the ack before + // the end of the kernel. There's plenty of time for the transactions to complete. + + // The 16B chunk written by the thread. + int co = tidx / QK_VECS_IN_16B; + // The position of the thread in that 16B chunk. + int ci = tidx % QK_VECS_IN_16B * QK_VEC_SIZE; + + // Two chunks are separated by L * x elements. A thread write QK_VEC_SIZE elements. + int offset = bhi_kv * params.memory_max_len * Dh + co * params.memory_max_len * QK_ELTS_IN_16B + + // params.timestep*QK_ELTS_IN_16B + + tlength_circ * QK_ELTS_IN_16B + ci; + + if (write_kv_cache) { + // Trigger the stores to global memory. + if (Dh == Dh_MAX || co < Dh / QK_ELTS_IN_16B) { + *reinterpret_cast(¶ms.k_cache[offset]) = k; + } + } + + // Compute \sum_i Q[i] * K^T[i] for the current timestep. +#ifdef MMHA_USE_FP32_ACUM_FOR_FMA + using Qk_vec_acum = typename Qk_vec_acum_fp32_::Type; +#else + using Qk_vec_acum = Qk_vec; +#endif + qk = dot(q, k); + if (QK_VECS_PER_WARP <= WARP_SIZE) { +#pragma unroll + for (int mask = QK_VECS_PER_WARP / 2; mask >= 1; mask /= 2) { + qk += __shfl_xor_sync(shfl_mask(QK_VECS_PER_WARP), qk, mask); + } + } + } + + if (QK_VECS_PER_WARP > WARP_SIZE) { + constexpr int WARPS_PER_RED = (QK_VECS_PER_WARP + WARP_SIZE - 1) / WARP_SIZE; + qk = block_sum(&red_smem[WARPS_PER_RED], qk); + } + + // Store that value in shared memory. Keep the Q*K^T value in register for softmax. + if (tidx == 0) { + // Normalize qk. + qk *= params.inv_sqrt_dh; + if (params.relative_attention_bias != nullptr) { + // TODO (Haotian): check whether we should replace hi with hi_kv, + // although params.relative_attention_bias is usually not used. + qk = add(qk, + params.relative_attention_bias[hi * params.relative_attention_bias_stride + * params.relative_attention_bias_stride + + (tlength - padd_len) * params.relative_attention_bias_stride + + (tlength - padd_len)]); + } + // Add alibi positional encoding + // qk += (alibi_slope != 0) ? alibi_slope * (params.timestep - params.memory_max_len) : 0; + // We don't need to apply the linear position bias here since qi - ki = 0 yields the position bias 0. + + qk_max = qk; + qk_smem[tlength - first_step] = qk; + // qk_smem[params.timestep] = qk; + } + + // Make sure the data is in shared memory. + __syncthreads(); + + // The type of queries and keys for the math in the Q*K^T product. + using K_vec = typename K_vec_::Type; + // The number of elements per vector. + constexpr int K_VEC_SIZE = sizeof(K_vec) / sizeof(T); + // Make sure the hidden size per head is a multiple of the vector size. + static_assert(Dh_MAX % K_VEC_SIZE == 0, ""); + // The number of elements per thread. + constexpr int K_ELTS_PER_THREAD = Dh_MAX / THREADS_PER_KEY; + // The number of vectors per thread. + constexpr int K_VECS_PER_THREAD = K_ELTS_PER_THREAD / K_VEC_SIZE; + + // The position the first key loaded by each thread from the cache buffer (for this B * H). + int ko = tidx / THREADS_PER_KEY; + // The position of the thread in the chunk of keys. + int ki = tidx % THREADS_PER_KEY * K_VEC_SIZE; + + static_assert(Dh_MAX == THREADS_PER_KEY * K_VEC_SIZE * K_VECS_PER_THREAD); + + // Load the Q values from shared memory. The values are reused during the loop on K. + K_vec q_vec[K_VECS_PER_THREAD]; +#pragma unroll + for (int ii = 0; ii < K_VECS_PER_THREAD; ++ii) { + q_vec[ii] = *reinterpret_cast(&q_smem[ki + ii * THREADS_PER_KEY * K_VEC_SIZE]); + } + + K_vec k_bias_vec[DO_CROSS_ATTENTION ? K_VECS_PER_THREAD : 1]; + if (DO_CROSS_ATTENTION && params.timestep == 0) { +#pragma unroll + for (int ii = 0; ii < K_VECS_PER_THREAD; ++ii) { + k_bias_vec[ii] = *reinterpret_cast(&bias_smem[ki + ii * THREADS_PER_KEY * K_VEC_SIZE]); + } + } + + // The number of timesteps loaded per iteration. + constexpr int K_PER_ITER = THREADS_PER_BLOCK / THREADS_PER_KEY; + // The number of keys per warp. + constexpr int K_PER_WARP = WARP_SIZE / THREADS_PER_KEY; + + // The base pointer for the key in the cache buffer. + T* k_cache = ¶ms.k_cache[bhi_kv * params.memory_max_len * Dh + ki]; + // Base pointer for the beam's batch, before offsetting with indirection buffer + T* k_cache_batch = ¶ms.k_cache[bbhi_kv * params.memory_max_len * Dh + ki]; + + // Pick a number of keys to make sure all the threads of a warp enter (due to shfl_sync). + // int ti_end = div_up(params.timestep, K_PER_WARP) * K_PER_WARP; + int ti_end = div_up(tlength - first_step, K_PER_WARP) * K_PER_WARP + first_step; + + // prefix prompt length if has + const int prefix_prompt_length = (params.prefix_prompt_lengths == nullptr) ? 0 : params.prefix_prompt_lengths[bi]; + + // Iterate over the keys/timesteps to compute the various (Q*K^T)_{ti} values. + const bool has_beams = params.cache_indir != nullptr; + const int* beam_indices = has_beams ? ¶ms.cache_indir[bi_seq_len_offset] : nullptr; + + for (int ti = first_step + ko; ti < ti_end; ti += K_PER_ITER) { + const int ti_circ = ti % params.memory_max_len; + + // The keys loaded from the key cache. + K_vec k[K_VECS_PER_THREAD]; + K_vec k_vec_zero; + zero(k_vec_zero); +#pragma unroll + for (int ii = 0; ii < K_VECS_PER_THREAD; ++ii) { + int jj = ii * params.memory_max_len + ti_circ; + // if( ti < params.timestep ) { + const bool within_bounds = (Dh == Dh_MAX || jj * QK_ELTS_IN_16B < Dh * params.memory_max_len); + if (ti < tlength) { + if (!within_bounds) { + k[ii] = k_vec_zero; + } + else { + if (has_beams) { + const int beam_offset = beam_indices[ti_circ] * params.num_heads * params.memory_max_len * Dh; + k[ii] = *reinterpret_cast(&k_cache_batch[beam_offset + jj * QK_ELTS_IN_16B]); + } + else { + k[ii] = *reinterpret_cast(&k_cache_batch[jj * QK_ELTS_IN_16B]); + } + } + // add bias and update k_cache + if (DO_CROSS_ATTENTION && params.timestep == 0) { + k[ii] = add(k[ii], k_bias_vec[ii]); + + if (do_ia3) { + k[ii] = mul( + k[ii], + *reinterpret_cast( + ¶ms.ia3_key_weights[(ia3_task_id * params.num_heads + hi) * Dh + ki + + ii * THREADS_PER_KEY * K_VEC_SIZE])); + } + + if (Dh == Dh_MAX || jj * QK_ELTS_IN_16B < Dh * params.memory_max_len) { + *reinterpret_cast(&k_cache[jj * QK_ELTS_IN_16B]) = k[ii]; + } + } + } + } + + // Perform the dot product and normalize qk. + // + // WARNING: ALL THE THREADS OF A WARP MUST ENTER!!! + float qk = Qk_dot::dot(q_vec, k) * params.inv_sqrt_dh; + bool is_mask = (params.masked_tokens != nullptr) && params.masked_tokens[bi_seq_len_offset + ti]; + + // Store the product to shared memory. There's one qk value per timestep. Update the max. + // if( ti < params.timestep && tidx % THREADS_PER_KEY == 0 ) { + if (ti < tlength && tidx % THREADS_PER_KEY == 0) { + if (params.relative_attention_bias != nullptr) { + qk = add(qk, + params.relative_attention_bias[hi * params.relative_attention_bias_stride + * params.relative_attention_bias_stride + + tlength * params.relative_attention_bias_stride + ti]); + } + if (params.linear_bias_slopes != nullptr) { + // Apply the linear position bias: (ki - qi) * slope[hi]. + // The padding token locates between the input context and the generated tokens. + // We need to remove the number of padding tokens in the distance computation. + // ti : 0 1 2 3 4 5 6 7 8 9(tlength) + // token: i i i i p p p o o o where i=input, p=pad, o=output. + // e.g. ti = 2, dist = (9 - 3) - 2 = 4. + int max_context_length = params.max_prefix_prompt_length + params.max_input_length; + float dist = (ti < max_context_length ? ti + padd_len : ti) - tlength; + + qk += mul(params.linear_bias_slopes[hi], dist); + } + // Add alibi positional encoding + // qk += (alibi_slope != 0) ? alibi_slope * (params.timestep - params.memory_max_len) : 0; + qk_max = is_mask ? qk_max : fmaxf(qk_max, qk); + qk_smem[ti - first_step] = qk; + } + } + +// Perform the final reduction to compute the max inside each warp. +// +// NOTE: In a group of THREADS_PER_KEY threads, the leader already has the max value for the +// group so it's not needed to run the reduction inside the group (again). +#pragma unroll + for (int mask = WARP_SIZE / 2; mask >= THREADS_PER_KEY; mask /= 2) { + qk_max = fmaxf(qk_max, __shfl_xor_sync(uint32_t(-1), qk_max, mask)); + } + + // Decompose the thread index into warp and lane. + const int warp = tidx / WARP_SIZE; + const int lane = tidx % WARP_SIZE; + + // The warp leader writes the max to shared memory. + if (lane == 0) { + red_smem[warp] = qk_max; + } + + // Make sure the products are in shared memory. + __syncthreads(); + + // The warps finalize the reduction. + qk_max = lane < WARPS_PER_BLOCK ? red_smem[lane] : -FLT_MAX; +#pragma unroll + for (int mask = WARPS_PER_BLOCK / 2; mask >= 1; mask /= 2) { + qk_max = fmaxf(qk_max, __shfl_xor_sync(uint32_t(-1), qk_max, mask)); + } + + // Broadcast to all the threads in the warp. + qk_max = __shfl_sync(uint32_t(-1), qk_max, 0); + + // Compute the logits and start the sum. + float sum = 0.f; + // for( int ti = tidx; ti <= params.timestep; ti += THREADS_PER_BLOCK ) { + for (int ti = first_step + tidx; ti <= tlength; ti += THREADS_PER_BLOCK) { + bool is_mask = (params.masked_tokens != nullptr) && params.masked_tokens[bi_seq_len_offset + ti]; + float logit = is_mask ? 0.f : __expf(qk_smem[ti - first_step] - qk_max); + sum += logit; + qk_smem[ti - first_step] = logit; + } + + // Compute the sum. + sum = block_sum(&red_smem[WARPS_PER_BLOCK], sum); + + // Normalize the logits. + float inv_sum = __fdividef(1.f, sum + 1.e-6f); + // for( int ti = tidx; ti <= params.timestep; ti += THREADS_PER_BLOCK ) { + const size_t cross_attention_out_offset = + params.is_return_cross_attentions ? + bhi_kv * params.max_decoder_seq_len * params.memory_max_len + params.timestep * params.memory_max_len : + 0; + for (int ti = first_step + tidx; ti <= tlength; ti += THREADS_PER_BLOCK) { + float logit = qk_smem[ti - first_step] * inv_sum; + if (params.is_return_cross_attentions) { + params.cross_attention_out[cross_attention_out_offset + ti] = logit; + } + convert_from_float(logits_smem[ti - first_step], logit); + } + + // Put Values part below so we leverage __syncthreads + // from the previous step + + // The number of elements per vector. + constexpr int V_VEC_SIZE = Dh_MAX / THREADS_PER_VALUE; + // A vector of V elements for the current timestep. + using V_vec = typename V_vec_::Type; + + // The value computed by this thread. + int vo = tidx / THREADS_PER_VALUE; + // The hidden dimensions computed by this particular thread. + int vi = tidx % THREADS_PER_VALUE * V_VEC_SIZE; + + // The base pointer for the value in the cache buffer. + T* v_cache = ¶ms.v_cache[bhi_kv * params.memory_max_len * Dh + vi]; + // Base pointer for the beam's batch, before offsetting with indirection buffer + T* v_cache_batch = ¶ms.v_cache[bbhi_kv * params.memory_max_len * Dh + vi]; + + // The number of values processed per iteration of the loop. + constexpr int V_PER_ITER = THREADS_PER_BLOCK / THREADS_PER_VALUE; + + // One group of threads computes the product(s) for the current timestep. + V_vec v_bias; + zero(v_bias); + // if( vo == params.timestep % V_PER_ITER ) { + if (Dh == Dh_MAX || vi < Dh) { + if (handle_kv) { + if (vo == tlength % V_PER_ITER) { + // Trigger the loads from the V bias buffer. + if (params.v_bias != nullptr) { + v_bias = *reinterpret_cast(¶ms.v_bias[hi_kv * Dh + vi]); + } + if (DO_CROSS_ATTENTION) { + *reinterpret_cast(&bias_smem[vi]) = v_bias; + } + } + } + } + + // From previous, before values, step + // Also make sure the logits are in shared memory. + __syncthreads(); + + // Values continued +#ifdef MMHA_USE_FP32_ACUM_FOR_OUT + using V_vec_acum = typename V_vec_acum_fp32_::Type; +#else + using V_vec_acum = V_vec; +#endif + // The partial outputs computed by each thread. + V_vec_acum out; + zero(out); + + // Loop over the timesteps to compute the partial outputs. + // for( int ti = vo; ti < params.timestep; ti += V_PER_ITER ) { + if (Dh == Dh_MAX || vi < Dh) { + for (int ti = first_step + vo; ti < tlength; ti += V_PER_ITER) { + const int ti_circ = ti % params.memory_max_len; + + // Fetch offset based on cache_indir when beam sampling + const int beam_src = (params.cache_indir != nullptr) ? params.cache_indir[bi_seq_len_offset + ti_circ] : 0; + const int beam_offset = beam_src * params.num_heads * params.memory_max_len * Dh; + // Load the values from the cache. + V_vec v = *reinterpret_cast(&v_cache_batch[beam_offset + ti_circ * Dh]); + if (DO_CROSS_ATTENTION && params.timestep == 0) { + v = add(v, *reinterpret_cast(&bias_smem[vi])); + if (do_ia3) { + v = mul( + v, + *reinterpret_cast( + ¶ms.ia3_value_weights[(ia3_task_id * params.num_heads + hi) * Dh + vi])); + } + *reinterpret_cast(&v_cache[ti * Dh]) = v; + } + // Load the logits from shared memory. +#if defined(MMHA_USE_FP32_ACUM_FOR_LOGITS) + float logit = logits_smem[ti - first_step]; + out = fma(logit, cast_to_float(v), out); +#else + T logit = logits_smem[ti - first_step]; + + // Update the partial sums. + out = fma(logit, v, out); +#endif + } + } + + // One group of threads computes the product(s) for the current timestep. + // if( vo == params.timestep % V_PER_ITER ) { + if (vo == tlength % V_PER_ITER && (Dh == Dh_MAX || vi < Dh)) { + + V_vec v; + if (DO_CROSS_ATTENTION) { + v = *reinterpret_cast(&v_cache[tlength * Dh]); + } + else { + // Trigger the loads from the V buffer. + const auto v_offset = v_base_offset + vi; + if (params.int8_mode == 2) { + using Packed_Int8_t = typename packed_type::value>::type; + using Packed_Float_t = typename packed_type::value>::type; + const auto v_scaling = params.qkv_scale_out[2]; + const auto v_quant = + *reinterpret_cast(&reinterpret_cast(params.v)[v_offset]); + + convert_from_float(v, mul(v_scaling, float_from_int8(v_quant))); + } + else { + v = *reinterpret_cast(¶ms.v[v_offset]); + } + // Trigger the loads from the V bias buffer. + // V_vec v_bias = *reinterpret_cast(¶ms.v_bias[hi*Dh + vi]); + } + + // Compute the V values with bias. + v = add(v, v_bias); + if (write_kv_cache) { + + if (do_ia3) { + v = mul( + v, + *reinterpret_cast( + ¶ms.ia3_value_weights[(ia3_task_id * params.num_heads + hi) * Dh + vi])); + } + + // Store the values with bias back to global memory in the cache for V. + //*reinterpret_cast(&v_cache[params.timestep*Dh]) = v; + *reinterpret_cast(&v_cache[tlength_circ * Dh]) = v; + } + + // Initialize the output value with the current timestep. +#if defined(MMHA_USE_FP32_ACUM_FOR_LOGITS) + // out = fma(logits_smem[params.timestep], cast_to_float(v), out); + out = fma(logits_smem[tlength - first_step], cast_to_float(v), out); +#else + // out = fma(logits_smem[params.timestep], v, out); + out = fma(logits_smem[tlength - first_step], v, out); +#endif + } + + // Make sure we can start writing to shared memory. + __syncthreads(); + + // Run the final reduction amongst the different groups computing different partial outputs. + if (Dh == Dh_MAX || vi < Dh) { +#pragma unroll + for (int active_groups = V_PER_ITER; active_groups >= 2; active_groups /= 2) { + + // The midpoint in the number of active groups. + int midpoint = active_groups / 2; + + // The upper part of active threads store to shared memory. + if (vo >= midpoint && vo < active_groups && (Dh == Dh_MAX || vi < Dh)) { +#ifdef MMHA_USE_FP32_ACUM_FOR_OUT + convert_from_float(*reinterpret_cast(&out_smem[(vo - midpoint) * Dh + vi]), out); +#else + *reinterpret_cast(&out_smem[(vo - midpoint) * Dh + vi]) = out; +#endif + } + __syncthreads(); + + // The bottom warps update their values. + if (vo < midpoint && (Dh == Dh_MAX || vi < Dh)) { + out = add(*reinterpret_cast(&out_smem[vo * Dh + vi]), out); + } + __syncthreads(); + } + } + + // Output the final values. + if (vo == 0 && (Dh == Dh_MAX || vi < Dh)) { +#ifdef MMHA_USE_FP32_ACUM_FOR_OUT + if (params.int8_mode == 2) { + using Packed_Int8_t = typename packed_type::value>::type; + out = mul(*params.attention_out_scale, out); + *reinterpret_cast(&(reinterpret_cast(params.out)[bhi * Dh + vi])) = + cast_to_int8(out); + } + else { + convert_from_float(*reinterpret_cast(¶ms.out[bhi * Dh + vi]), out); + } +#else + // TODO: support int8_mode? + *reinterpret_cast(¶ms.out[bhi * Dh + vi]) = out; +#endif + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mmha + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +void mmha_launch_kernel(const KERNEL_PARAMS_TYPE& params, const cudaStream_t& stream); diff --git a/llm-awq/awq/kernels/csrc/attention/decoder_masked_multihead_attention_utils.h b/llm-awq/awq/kernels/csrc/attention/decoder_masked_multihead_attention_utils.h new file mode 100644 index 0000000000000000000000000000000000000000..7f6b7d81f69204c0f16cd5c2647196486c66ae08 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/attention/decoder_masked_multihead_attention_utils.h @@ -0,0 +1,1795 @@ +// Downloaded from from FasterTransformer v5.2.1 +// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.2.1_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention_utils.h +/* + * Copyright (c) 2020-2022, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "cuda_bf16_wrapper.h" +#include "cuda_bf16_fallbacks.cuh" +#include + +using namespace fastertransformer; + +namespace mmha { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +struct Float8_ { + float2 x; + float2 y; + float2 z; + float2 w; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +struct Float4_ { + float2 x; + float2 y; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +#ifdef ENABLE_BF16 +struct bf16_4_t { + __nv_bfloat162 x; + __nv_bfloat162 y; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +struct bf16_8_t { + __nv_bfloat162 x; + __nv_bfloat162 y; + __nv_bfloat162 z; + __nv_bfloat162 w; +}; +#endif + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct num_elems; +template<> +struct num_elems { + static constexpr int value = 1; +}; +template<> +struct num_elems { + static constexpr int value = 2; +}; +template<> +struct num_elems { + static constexpr int value = 4; +}; +template<> +struct num_elems { + static constexpr int value = 4; +}; +template<> +struct num_elems { + static constexpr int value = 8; +}; + +template<> +struct num_elems { + static constexpr int value = 2; +}; +template<> +struct num_elems { + static constexpr int value = 4; +}; +template<> +struct num_elems { + static constexpr int value = 8; +}; + +#ifdef ENABLE_BF16 +template<> +struct num_elems<__nv_bfloat162> { + static constexpr int value = 2; +}; +template<> +struct num_elems { + static constexpr int value = 4; +}; +template<> +struct num_elems { + static constexpr int value = 8; +}; +#endif + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct packed_type; +template +struct packed_type { + using type = T; +}; +template<> +struct packed_type { + using type = int16_t; +}; +template<> +struct packed_type { + using type = int32_t; +}; +template<> +struct packed_type { + using type = int64_t; +}; + +template<> +struct packed_type { + using type = float2; +}; +template<> +struct packed_type { + using type = float4; +}; +template<> +struct packed_type { + using type = Float8_; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float add(float a, float b) +{ + return a + b; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float2 add(float2 a, float2 b) +{ + float2 c; + c.x = add(a.x, b.x); + c.y = add(a.y, b.y); + return c; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float4 add(float4 a, float4 b) +{ + float4 c; + c.x = add(a.x, b.x); + c.y = add(a.y, b.y); + c.z = add(a.z, b.z); + c.w = add(a.w, b.w); + return c; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +#ifdef ENABLE_BF16 +inline __device__ __nv_bfloat16 add(__nv_bfloat16 a, __nv_bfloat16 b) +{ + return a + b; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ __nv_bfloat162 add(__nv_bfloat162 a, __nv_bfloat162 b) +{ + return bf16hadd2(a, b); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ bf16_4_t add(bf16_4_t a, bf16_4_t b) +{ + bf16_4_t c; + c.x = add(a.x, b.x); + c.y = add(a.y, b.y); + return c; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ bf16_8_t add(bf16_8_t a, bf16_8_t b) +{ + bf16_8_t c; + c.x = add(a.x, b.x); + c.y = add(a.y, b.y); + c.z = add(a.z, b.z); + c.w = add(a.w, b.w); + return c; +} +#endif // ENABLE_BF16 + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ uint16_t add(uint16_t a, uint16_t b) +{ + uint16_t c; + asm volatile("add.f16 %0, %1, %2;\n" : "=h"(c) : "h"(a), "h"(b)); + return c; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ uint32_t add(uint32_t a, uint32_t b) +{ + uint32_t c; + asm volatile("add.f16x2 %0, %1, %2;\n" : "=r"(c) : "r"(a), "r"(b)); + return c; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ uint2 add(uint2 a, uint2 b) +{ + uint2 c; + c.x = add(a.x, b.x); + c.y = add(a.y, b.y); + return c; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ uint4 add(uint4 a, uint4 b) +{ + uint4 c; + c.x = add(a.x, b.x); + c.y = add(a.y, b.y); + c.z = add(a.z, b.z); + c.w = add(a.w, b.w); + return c; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ uint16_t float_to_half(float f) +{ + union { + uint32_t u32; + uint16_t u16[2]; + } tmp; +#if 0 && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 // Is it better? + float zero = 0.f; + asm volatile("cvt.rn.f16x2.f32 %0, %1, %2;\n" : "=r"(tmp.u32) : "f"(zero), "f"(f)); +#else + asm volatile("cvt.rn.f16.f32 %0, %1;\n" : "=h"(tmp.u16[0]) : "f"(f)); +#endif + return tmp.u16[0]; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ uint32_t float2_to_half2(float2 f) +{ + union { + uint32_t u32; + uint16_t u16[2]; + } tmp; +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + asm volatile("cvt.rn.f16x2.f32 %0, %1, %2;\n" : "=r"(tmp.u32) : "f"(f.y), "f"(f.x)); +#else + asm volatile("cvt.rn.f16.f32 %0, %1;\n" : "=h"(tmp.u16[0]) : "f"(f.x)); + asm volatile("cvt.rn.f16.f32 %0, %1;\n" : "=h"(tmp.u16[1]) : "f"(f.y)); +#endif + return tmp.u32; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float half_to_float(uint16_t h) +{ + float f; + asm volatile("cvt.f32.f16 %0, %1;\n" : "=f"(f) : "h"(h)); + return f; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float2 half2_to_float2(uint32_t v) +{ + uint16_t lo, hi; + asm volatile("mov.b32 {%0, %1}, %2;\n" : "=h"(lo), "=h"(hi) : "r"(v)); + return make_float2(half_to_float(lo), half_to_float(hi)); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float add(float a, uint16_t b) +{ + return a + half_to_float(b); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +#ifdef ENABLE_BF16 +inline __device__ float add(float a, __nv_bfloat16 b) +{ + return a + __bfloat162float(b); +} +#endif + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float2 add(uint32_t a, float2 fb) +{ + float2 fa = half2_to_float2(a); + return add(fa, fb); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ Float4_ add(uint2 a, Float4_ fb) +{ + Float4_ fc; + fc.x = add(a.x, fb.x); + fc.y = add(a.y, fb.y); + return fc; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ Float8_ add(uint4 a, Float8_ fb) +{ + Float8_ fc; + fc.x = add(a.x, fb.x); + fc.y = add(a.y, fb.y); + fc.z = add(a.z, fb.z); + fc.w = add(a.w, fb.w); + return fc; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ uint32_t h0_h0(uint16_t a) +{ + uint32_t b; + asm volatile("mov.b32 %0, {%1, %1};" : "=r"(b) : "h"(a)); + return b; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float fma(float a, float b, float c) +{ + return a * b + c; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float2 fma(float2 a, float2 b, float2 c) +{ + float2 d; + d.x = fma(a.x, b.x, c.x); + d.y = fma(a.y, b.y, c.y); + return d; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float2 fma(float a, float2 b, float2 c) +{ + float2 d; + d.x = fma(a, b.x, c.x); + d.y = fma(a, b.y, c.y); + return d; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float4 fma(float4 a, float4 b, float4 c) +{ + float4 d; + d.x = fma(a.x, b.x, c.x); + d.y = fma(a.y, b.y, c.y); + d.z = fma(a.z, b.z, c.z); + d.w = fma(a.w, b.w, c.w); + return d; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float4 fma(float a, float4 b, float4 c) +{ + float4 d; + d.x = fma(a, b.x, c.x); + d.y = fma(a, b.y, c.y); + d.z = fma(a, b.z, c.z); + d.w = fma(a, b.w, c.w); + return d; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ Float4_ fma(float a, Float4_ b, Float4_ c) +{ + Float4_ d; + d.x = fma(a, b.x, c.x); + d.y = fma(a, b.y, c.y); + return d; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ Float8_ fma(float a, Float8_ b, Float8_ c) +{ + Float8_ d; + d.x = fma(a, b.x, c.x); + d.y = fma(a, b.y, c.y); + d.z = fma(a, b.z, c.z); + d.w = fma(a, b.w, c.w); + return d; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +#ifdef ENABLE_BF16 +inline __device__ float2 add(__nv_bfloat162 a, float2 fb) +{ + float2 fa = bf1622float2(a); + return add(fa, fb); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ Float4_ add(bf16_4_t a, Float4_ fb) +{ + Float4_ fc; + fc.x = add(a.x, fb.x); + fc.y = add(a.y, fb.y); + return fc; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ Float8_ add(bf16_8_t a, Float8_ fb) +{ + Float8_ fc; + fc.x = add(a.x, fb.x); + fc.y = add(a.y, fb.y); + fc.z = add(a.z, fb.z); + fc.w = add(a.w, fb.w); + return fc; +} +#endif // ENABLE_BF16 + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ uint32_t fma(uint32_t a, uint32_t b, uint32_t c) +{ + uint32_t d; + asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n" : "=r"(d) : "r"(a), "r"(b), "r"(c)); + return d; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ uint32_t fma(uint16_t a, uint32_t b, uint32_t c) +{ + return fma(h0_h0(a), b, c); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ uint2 fma(uint2 a, uint2 b, uint2 c) +{ + uint2 d; + d.x = fma(a.x, b.x, c.x); + d.y = fma(a.y, b.y, c.y); + return d; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ uint2 fma(uint16_t a, uint2 b, uint2 c) +{ + uint32_t s = h0_h0(a); + uint2 d; + d.x = fma(s, b.x, c.x); + d.y = fma(s, b.y, c.y); + return d; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ uint4 fma(uint4 a, uint4 b, uint4 c) +{ + uint4 d; + d.x = fma(a.x, b.x, c.x); + d.y = fma(a.y, b.y, c.y); + d.z = fma(a.z, b.z, c.z); + d.w = fma(a.w, b.w, c.w); + return d; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ uint4 fma(uint16_t a, uint4 b, uint4 c) +{ + uint32_t s = h0_h0(a); + uint4 d; + d.x = fma(s, b.x, c.x); + d.y = fma(s, b.y, c.y); + d.z = fma(s, b.z, c.z); + d.w = fma(s, b.w, c.w); + return d; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float fma(uint16_t a, uint16_t b, float fc) +{ + float fa = half_to_float(a); + float fb = half_to_float(b); + return fa * fb + fc; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float2 fma(uint32_t a, uint32_t b, float2 fc) +{ + float2 fa = half2_to_float2(a); + float2 fb = half2_to_float2(b); + return fma(fa, fb, fc); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float2 fma(uint16_t a, uint32_t b, float2 fc) +{ + return fma(h0_h0(a), b, fc); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ Float4_ fma(uint2 a, uint2 b, Float4_ fc) +{ + Float4_ fd; + fd.x = fma(a.x, b.x, fc.x); + fd.y = fma(a.y, b.y, fc.y); + return fd; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ Float4_ fma(uint16_t a, uint2 b, Float4_ fc) +{ + uint32_t s = h0_h0(a); + Float4_ fd; + fd.x = fma(s, b.x, fc.x); + fd.y = fma(s, b.y, fc.y); + return fd; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ Float8_ fma(uint4 a, uint4 b, Float8_ fc) +{ + Float8_ fd; + fd.x = fma(a.x, b.x, fc.x); + fd.y = fma(a.y, b.y, fc.y); + fd.z = fma(a.z, b.z, fc.z); + fd.w = fma(a.w, b.w, fc.w); + return fd; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ Float8_ fma(uint16_t a, uint4 b, Float8_ fc) +{ + uint32_t s = h0_h0(a); + Float8_ fd; + fd.x = fma(s, b.x, fc.x); + fd.y = fma(s, b.y, fc.y); + fd.z = fma(s, b.z, fc.z); + fd.w = fma(s, b.w, fc.w); + return fd; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// +#ifdef ENABLE_BF16 +inline __device__ __nv_bfloat162 fma(__nv_bfloat162 a, __nv_bfloat162 b, __nv_bfloat162 c) +{ + return bf16hfma2(a, b, c); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ __nv_bfloat162 fma(__nv_bfloat16 a, __nv_bfloat162 b, __nv_bfloat162 c) +{ + return bf16hfma2(bf162bf162(a), b, c); +} +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ bf16_4_t fma(bf16_4_t a, bf16_4_t b, bf16_4_t c) +{ + bf16_4_t d; + d.x = fma(a.x, b.x, c.x); + d.y = fma(a.y, b.y, c.y); + return d; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ bf16_4_t fma(__nv_bfloat16 a, bf16_4_t b, bf16_4_t c) +{ + __nv_bfloat162 s = bf162bf162(a); + bf16_4_t d; + d.x = fma(s, b.x, c.x); + d.y = fma(s, b.y, c.y); + return d; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ bf16_8_t fma(bf16_8_t a, bf16_8_t b, bf16_8_t c) +{ + bf16_8_t d; + d.x = fma(a.x, b.x, c.x); + d.y = fma(a.y, b.y, c.y); + d.z = fma(a.z, b.z, c.z); + d.w = fma(a.w, b.w, c.w); + return d; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ bf16_8_t fma(__nv_bfloat16 a, bf16_8_t b, bf16_8_t c) +{ + __nv_bfloat162 s = bf162bf162(a); + bf16_8_t d; + d.x = fma(s, b.x, c.x); + d.y = fma(s, b.y, c.y); + d.z = fma(s, b.z, c.z); + d.w = fma(s, b.w, c.w); + return d; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float fma(__nv_bfloat16 a, __nv_bfloat16 b, float fc) +{ + return __bfloat162float(a) * __bfloat162float(b) + fc; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float2 fma(__nv_bfloat162 a, __nv_bfloat162 b, float2 fc) +{ + float2 fa = bf1622float2(a); + float2 fb = bf1622float2(b); + return fma(fa, fb, fc); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float2 fma(__nv_bfloat16 a, __nv_bfloat162 b, float2 fc) +{ + return fma(bf162bf162(a), b, fc); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ Float4_ fma(bf16_4_t a, bf16_4_t b, Float4_ fc) +{ + Float4_ fd; + fd.x = fma(a.x, b.x, fc.x); + fd.y = fma(a.y, b.y, fc.y); + return fd; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ Float4_ fma(__nv_bfloat16 a, bf16_4_t b, Float4_ fc) +{ + __nv_bfloat162 s = bf162bf162(a); + Float4_ fd; + fd.x = fma(s, b.x, fc.x); + fd.y = fma(s, b.y, fc.y); + return fd; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ Float8_ fma(bf16_8_t a, bf16_8_t b, Float8_ fc) +{ + Float8_ fd; + fd.x = fma(a.x, b.x, fc.x); + fd.y = fma(a.y, b.y, fc.y); + fd.z = fma(a.z, b.z, fc.z); + fd.w = fma(a.w, b.w, fc.w); + return fd; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ Float8_ fma(__nv_bfloat16 a, bf16_8_t b, Float8_ fc) +{ + __nv_bfloat162 s = bf162bf162(a); + Float8_ fd; + fd.x = fma(s, b.x, fc.x); + fd.y = fma(s, b.y, fc.y); + fd.z = fma(s, b.z, fc.z); + fd.w = fma(s, b.w, fc.w); + return fd; +} +#endif // ENABLE_BF16 +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +inline __device__ Acc mul(A a, B b) +{ + return a * b; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template<> +inline __device__ float mul(float a, float b) +{ + return a * b; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template<> +inline __device__ float2 mul(float2 a, float2 b) +{ + float2 c; + c.x = a.x * b.x; + c.y = a.y * b.y; + return c; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template<> +inline __device__ float2 mul(float a, float2 b) +{ + float2 c; + c.x = a * b.x; + c.y = a * b.y; + return c; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template<> +inline __device__ float4 mul(float4 a, float4 b) +{ + float4 c; + c.x = a.x * b.x; + c.y = a.y * b.y; + c.z = a.z * b.z; + c.w = a.w * b.w; + return c; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template<> +inline __device__ float4 mul(float a, float4 b) +{ + float4 c; + c.x = a * b.x; + c.y = a * b.y; + c.z = a * b.z; + c.w = a * b.w; + return c; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template<> +inline __device__ Float8_ mul(float a, Float8_ b) +{ + Float8_ c; + c.x = make_float2(a * b.x.x, a * b.x.y); + c.y = make_float2(a * b.y.x, a * b.y.y); + c.z = make_float2(a * b.z.x, a * b.z.y); + c.w = make_float2(a * b.w.x, a * b.w.y); + return c; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template<> +inline __device__ uint16_t mul(uint16_t a, uint16_t b) +{ + uint16_t c; + asm volatile("mul.f16 %0, %1, %2;\n" : "=h"(c) : "h"(a), "h"(b)); + return c; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template<> +inline __device__ uint32_t mul(uint32_t a, uint32_t b) +{ + uint32_t c; + asm volatile("mul.f16x2 %0, %1, %2;\n" : "=r"(c) : "r"(a), "r"(b)); + return c; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template<> +inline __device__ uint32_t mul(uint16_t a, uint32_t b) +{ + return mul(h0_h0(a), b); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template<> +inline __device__ uint2 mul(uint2 a, uint2 b) +{ + uint2 c; + c.x = mul(a.x, b.x); + c.y = mul(a.y, b.y); + return c; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template<> +inline __device__ uint2 mul(uint16_t a, uint2 b) +{ + uint32_t s = h0_h0(a); + uint2 c; + c.x = mul(s, b.x); + c.y = mul(s, b.y); + return c; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template<> +inline __device__ uint4 mul(uint4 a, uint4 b) +{ + uint4 c; + c.x = mul(a.x, b.x); + c.y = mul(a.y, b.y); + c.z = mul(a.z, b.z); + c.w = mul(a.w, b.w); + return c; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template<> +inline __device__ uint4 mul(uint16_t a, uint4 b) +{ + uint32_t s = h0_h0(a); + uint4 c; + c.x = mul(s, b.x); + c.y = mul(s, b.y); + c.z = mul(s, b.z); + c.w = mul(s, b.w); + return c; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template<> +inline __device__ float mul(uint16_t a, uint16_t b) +{ + float fa = half_to_float(a); + float fb = half_to_float(b); + return fa * fb; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template<> +inline __device__ float mul(uint16_t a, float b) +{ + return half_to_float(a) * b; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template<> +inline __device__ float2 mul(uint32_t a, uint32_t b) +{ + float2 fa = half2_to_float2(a); + float2 fb = half2_to_float2(b); + return mul(fa, fb); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template<> +inline __device__ float2 mul(uint16_t a, uint32_t b) +{ + return mul(h0_h0(a), b); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template<> +inline __device__ Float4_ mul(uint2 a, uint2 b) +{ + Float4_ fc; + fc.x = mul(a.x, b.x); + fc.y = mul(a.y, b.y); + return fc; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template<> +inline __device__ Float4_ mul(uint16_t a, uint2 b) +{ + uint32_t s = h0_h0(a); + Float4_ fc; + fc.x = mul(s, b.x); + fc.y = mul(s, b.y); + return fc; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template<> +inline __device__ Float8_ mul(uint4 a, uint4 b) +{ + Float8_ fc; + fc.x = mul(a.x, b.x); + fc.y = mul(a.y, b.y); + fc.z = mul(a.z, b.z); + fc.w = mul(a.w, b.w); + return fc; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template<> +inline __device__ Float8_ mul(uint16_t a, uint4 b) +{ + uint32_t s = h0_h0(a); + Float8_ fc; + fc.x = mul(s, b.x); + fc.y = mul(s, b.y); + fc.z = mul(s, b.z); + fc.w = mul(s, b.w); + return fc; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +#ifdef ENABLE_BF16 +template<> +inline __device__ __nv_bfloat16 mul(__nv_bfloat16 a, __nv_bfloat16 b) +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + return __hmul(a, b); +#else + return bf16hmul(a, b); +#endif +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template<> +inline __device__ __nv_bfloat162 mul(__nv_bfloat162 a, __nv_bfloat162 b) +{ + return bf16hmul2(a, b); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template<> +inline __device__ __nv_bfloat162 mul(__nv_bfloat16 a, __nv_bfloat162 b) +{ + return mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(bf162bf162(a), b); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template<> +inline __device__ bf16_4_t mul(bf16_4_t a, bf16_4_t b) +{ + bf16_4_t c; + c.x = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(a.x, b.x); + c.y = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(a.y, b.y); + return c; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template<> +inline __device__ bf16_4_t mul(__nv_bfloat16 a, bf16_4_t b) +{ + __nv_bfloat162 s = bf162bf162(a); + bf16_4_t c; + c.x = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(s, b.x); + c.y = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(s, b.y); + return c; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template<> +inline __device__ bf16_8_t mul(bf16_8_t a, bf16_8_t b) +{ + bf16_8_t c; + c.x = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(a.x, b.x); + c.y = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(a.y, b.y); + c.z = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(a.z, b.z); + c.w = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(a.w, b.w); + return c; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template<> +inline __device__ bf16_8_t mul(__nv_bfloat16 a, bf16_8_t b) +{ + __nv_bfloat162 s = bf162bf162(a); + bf16_8_t c; + c.x = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(s, b.x); + c.y = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(s, b.y); + c.z = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(s, b.z); + c.w = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(s, b.w); + return c; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template<> +inline __device__ float mul(__nv_bfloat16 a, __nv_bfloat16 b) +{ + float fa = (float)a; + float fb = (float)b; + return fa * fb; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template<> +inline __device__ float mul(__nv_bfloat16 a, float b) +{ + return __bfloat162float(a) * b; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template<> +inline __device__ float2 mul(__nv_bfloat162 a, __nv_bfloat162 b) +{ + float2 fa = bf1622float2(a); + float2 fb = bf1622float2(b); + return mul(fa, fb); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template<> +inline __device__ float2 mul(__nv_bfloat16 a, __nv_bfloat162 b) +{ + return mul(bf162bf162(a), b); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template<> +inline __device__ Float4_ mul(bf16_4_t a, bf16_4_t b) +{ + Float4_ fc; + fc.x = mul(a.x, b.x); + fc.y = mul(a.y, b.y); + return fc; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template<> +inline __device__ Float4_ mul(__nv_bfloat16 a, bf16_4_t b) +{ + __nv_bfloat162 s = bf162bf162(a); + Float4_ fc; + fc.x = mul(s, b.x); + fc.y = mul(s, b.y); + return fc; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template<> +inline __device__ Float8_ mul(bf16_8_t a, bf16_8_t b) +{ + Float8_ fc; + fc.x = mul(a.x, b.x); + fc.y = mul(a.y, b.y); + fc.z = mul(a.z, b.z); + fc.w = mul(a.w, b.w); + return fc; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template<> +inline __device__ Float8_ mul(__nv_bfloat16 a, bf16_8_t b) +{ + __nv_bfloat162 s = bf162bf162(a); + Float8_ fc; + fc.x = mul(s, b.x); + fc.y = mul(s, b.y); + fc.z = mul(s, b.z); + fc.w = mul(s, b.w); + return fc; +} +#endif // ENABLE_BF16 +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float sum(float v) +{ + return v; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float sum(float2 v) +{ + return v.x + v.y; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float sum(float4 v) +{ + return v.x + v.y + v.z + v.w; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +#ifdef ENABLE_BF16 +inline __device__ float sum(__nv_bfloat162 v) +{ + float2 vf = bf1622float2(v); + return vf.x + vf.y; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float sum(bf16_4_t v) +{ + return sum(v.x) + sum(v.y); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float sum(bf16_8_t v) +{ + return sum(v.x) + sum(v.y) + sum(v.z) + sum(v.w); +} +#endif // ENABLE_BF16 +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float sum(uint16_t v) +{ + return half_to_float(v); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float sum(uint32_t v) +{ + float2 tmp = half2_to_float2(v); + return tmp.x + tmp.y; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float sum(uint2 v) +{ + uint32_t c = add(v.x, v.y); + return sum(c); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float sum(uint4 v) +{ +#if 1 + uint32_t c = add(v.x, v.y); + c = add(c, v.z); + c = add(c, v.w); +#else + uint32_t c = add(v.x, v.y); + uint32_t d = add(v.z, v.w); + c = add(c, d); +#endif + return sum(c); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float sum(Float4_ v) +{ + return v.x.x + v.x.y + v.y.x + v.y.y; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float sum(Float8_ v) +{ + return v.x.x + v.x.y + v.y.x + v.y.y + v.z.x + v.z.y + v.w.x + v.w.y; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +inline __device__ float dot(T a, T b) +{ + return sum(mul(a, b)); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +inline __device__ float dot(T a, T b) +{ + return sum(mul(a, b)); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ void zero(uint16_t& dst) +{ + dst = uint16_t(0); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +inline __device__ void zero(T& dst) +{ + constexpr int WORDS = sizeof(T) / 4; + union { + T raw; + uint32_t words[WORDS]; + } tmp; +#pragma unroll + for (int ii = 0; ii < WORDS; ++ii) { + tmp.words[ii] = 0u; + } + dst = tmp.raw; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// inline __device__ float2 rotary_embedding_coefficient(const int zid, const int rot_embed_dim, const float t_step, const float base) +// { +// const float inv_freq = t_step / pow(base, zid / (float)rot_embed_dim); +// return {cos(inv_freq), sin(inv_freq)}; +// } + +// with scale +inline __device__ float2 rotary_embedding_coefficient( + const int zid, const int rot_embed_dim, const float t_step, const float base, const float scale) +{ + const float inv_freq = (t_step * scale) / pow(base, zid / (float)rot_embed_dim); + return {cos(inv_freq), sin(inv_freq)}; +} + + +inline __device__ float2 rotary_embedding_transform(const float2 v, const float2 coef) +{ + float2 rot_v; + rot_v.x = coef.x * v.x - coef.y * v.y; + rot_v.y = coef.x * v.y + coef.y * v.x; + return rot_v; +} + +inline __device__ uint32_t rotary_embedding_transform(const uint32_t v, const float2 coef) +{ + float2 fv = half2_to_float2(v); + float2 rot_fv = rotary_embedding_transform(fv, coef); + return float2_to_half2(rot_fv); +} + +#ifdef ENABLE_BF16 +inline __device__ __nv_bfloat162 rotary_embedding_transform(const __nv_bfloat162 v, const float2 coef) +{ + float2 fv = bf1622float2(v); + float2 rot_fv = rotary_embedding_transform(fv, coef); + return __floats2bfloat162_rn(rot_fv.x, rot_fv.y); +} +#endif + +inline __device__ void apply_rotary_embedding(float& q, int zid, int rot_embed_dim, int t_step, const float base=10000.0f, const float scale=1.0f) +{ + return; +} + +inline __device__ void apply_rotary_embedding(float& q, float& k, int zid, int rot_embed_dim, int t_step, const float base=10000.0f, const float scale=1.0f) +{ + return; +} + +inline __device__ void apply_rotary_embedding(float2& q, int tid, int rot_embed_dim, int t_step, const float base=10000.0f, const float scale=1.0f) +{ + if (2 * tid >= rot_embed_dim) { + return; + } + const auto coef = rotary_embedding_coefficient(2 * tid, rot_embed_dim, t_step, base, scale); + q = rotary_embedding_transform(q, coef); +} + +inline __device__ void apply_rotary_embedding(float2& q, float2& k, int tid, int rot_embed_dim, int t_step, const float base=10000.0f, const float scale=1.0f) +{ + if (2 * tid >= rot_embed_dim) { + return; + } + const auto coef = rotary_embedding_coefficient(2 * tid, rot_embed_dim, t_step, base, scale); + q = rotary_embedding_transform(q, coef); + k = rotary_embedding_transform(k, coef); +} + +inline __device__ void apply_rotary_embedding(float4& q, int tid, int rot_embed_dim, int t_step, const float base=10000.0f, const float scale=1.0f) +{ + if (4 * tid >= rot_embed_dim) { + return; + } + + Float4_& q_ = *reinterpret_cast(&q); + const auto coef0 = rotary_embedding_coefficient(4 * tid, rot_embed_dim, t_step, base, scale); + q_.x = rotary_embedding_transform(q_.x, coef0); + const auto coef1 = rotary_embedding_coefficient(4 * tid + 2, rot_embed_dim, t_step, base, scale); + q_.y = rotary_embedding_transform(q_.y, coef1); +} + +inline __device__ void apply_rotary_embedding(float4& q, float4& k, int tid, int rot_embed_dim, int t_step, const float base=10000.0f, const float scale=1.0f) +{ + if (4 * tid >= rot_embed_dim) { + return; + } + + Float4_& q_ = *reinterpret_cast(&q); + Float4_& k_ = *reinterpret_cast(&k); + const auto coef0 = rotary_embedding_coefficient(4 * tid, rot_embed_dim, t_step, base, scale); + q_.x = rotary_embedding_transform(q_.x, coef0); + k_.x = rotary_embedding_transform(k_.x, coef0); + const auto coef1 = rotary_embedding_coefficient(4 * tid + 2, rot_embed_dim, t_step, base, scale); + q_.y = rotary_embedding_transform(q_.y, coef1); + k_.y = rotary_embedding_transform(k_.y, coef1); +} + +inline __device__ void apply_rotary_embedding(uint32_t& q, int tid, int rot_embed_dim, int t_step, const float base=10000.0f, const float scale=1.0f) +{ + if (2 * tid >= rot_embed_dim) { + return; + } + const auto coef = rotary_embedding_coefficient(2 * tid, rot_embed_dim, t_step, base, scale); + q = rotary_embedding_transform(q, coef); +} + +inline __device__ void apply_rotary_embedding(uint32_t& q, uint32_t& k, int tid, int rot_embed_dim, int t_step, const float base=10000.0f, const float scale=1.0f) +{ + if (2 * tid >= rot_embed_dim) { + return; + } + const auto coef = rotary_embedding_coefficient(2 * tid, rot_embed_dim, t_step, base, scale); + q = rotary_embedding_transform(q, coef); + k = rotary_embedding_transform(k, coef); +} + +inline __device__ void apply_rotary_embedding(uint2& q, int tid, int rot_embed_dim, int t_step, const float base=10000.0f, const float scale=1.0f) +{ + if (4 * tid >= rot_embed_dim) { + return; + } + const auto coef0 = rotary_embedding_coefficient(4 * tid, rot_embed_dim, t_step, base, scale); + q.x = rotary_embedding_transform(q.x, coef0); + const auto coef1 = rotary_embedding_coefficient(4 * tid + 2, rot_embed_dim, t_step, base, scale); + q.y = rotary_embedding_transform(q.y, coef1); +} + +inline __device__ void apply_rotary_embedding(uint2& q, uint2& k, int tid, int rot_embed_dim, int t_step, const float base=10000.0f, const float scale=1.0f) +{ + if (4 * tid >= rot_embed_dim) { + return; + } + const auto coef0 = rotary_embedding_coefficient(4 * tid, rot_embed_dim, t_step, base, scale); + q.x = rotary_embedding_transform(q.x, coef0); + k.x = rotary_embedding_transform(k.x, coef0); + const auto coef1 = rotary_embedding_coefficient(4 * tid + 2, rot_embed_dim, t_step, base, scale); + q.y = rotary_embedding_transform(q.y, coef1); + k.y = rotary_embedding_transform(k.y, coef1); +} + +inline __device__ void apply_rotary_embedding(uint4& q, int tid, int rot_embed_dim, int t_step, const float base=10000.0f, const float scale=1.0f) +{ + if (8 * tid >= rot_embed_dim) { + return; + } + const auto coef0 = rotary_embedding_coefficient(8 * tid, rot_embed_dim, t_step, base, scale); + q.x = rotary_embedding_transform(q.x, coef0); + const auto coef1 = rotary_embedding_coefficient(8 * tid + 2, rot_embed_dim, t_step, base, scale); + q.y = rotary_embedding_transform(q.y, coef1); + const auto coef2 = rotary_embedding_coefficient(8 * tid + 4, rot_embed_dim, t_step, base, scale); + q.z = rotary_embedding_transform(q.z, coef2); + const auto coef3 = rotary_embedding_coefficient(8 * tid + 6, rot_embed_dim, t_step, base, scale); + q.w = rotary_embedding_transform(q.w, coef3); +} + +inline __device__ void apply_rotary_embedding(uint4& q, uint4& k, int tid, int rot_embed_dim, int t_step, const float base=10000.0f, const float scale=1.0f) +{ + if (8 * tid >= rot_embed_dim) { + return; + } + const auto coef0 = rotary_embedding_coefficient(8 * tid, rot_embed_dim, t_step, base, scale); + q.x = rotary_embedding_transform(q.x, coef0); + k.x = rotary_embedding_transform(k.x, coef0); + const auto coef1 = rotary_embedding_coefficient(8 * tid + 2, rot_embed_dim, t_step, base, scale); + q.y = rotary_embedding_transform(q.y, coef1); + k.y = rotary_embedding_transform(k.y, coef1); + const auto coef2 = rotary_embedding_coefficient(8 * tid + 4, rot_embed_dim, t_step, base, scale); + q.z = rotary_embedding_transform(q.z, coef2); + k.z = rotary_embedding_transform(k.z, coef2); + const auto coef3 = rotary_embedding_coefficient(8 * tid + 6, rot_embed_dim, t_step, base, scale); + q.w = rotary_embedding_transform(q.w, coef3); + k.w = rotary_embedding_transform(k.w, coef3); +} + +#ifdef ENABLE_BF16 +inline __device__ void apply_rotary_embedding(__nv_bfloat162& q, int tid, int rot_embed_dim, int t_step, const float base=10000.0f, const float scale=1.0f) +{ + if (2 * tid >= rot_embed_dim) { + return; + } + const auto coef = rotary_embedding_coefficient(2 * tid, rot_embed_dim, t_step, base, scale); + q = rotary_embedding_transform(q, coef); +} + +inline __device__ void +apply_rotary_embedding(__nv_bfloat162& q, __nv_bfloat162& k, int tid, int rot_embed_dim, int t_step, const float base=10000.0f, const float scale=1.0f) +{ + if (2 * tid >= rot_embed_dim) { + return; + } + const auto coef = rotary_embedding_coefficient(2 * tid, rot_embed_dim, t_step, base, scale); + q = rotary_embedding_transform(q, coef); + k = rotary_embedding_transform(k, coef); +} + +inline __device__ void apply_rotary_embedding(bf16_4_t& q, int tid, int rot_embed_dim, int t_step, const float base=10000.0f, const float scale=1.0f) +{ + if (4 * tid >= rot_embed_dim) { + return; + } + const auto coef0 = rotary_embedding_coefficient(4 * tid, rot_embed_dim, t_step, base, scale); + q.x = rotary_embedding_transform(q.x, coef0); + const auto coef1 = rotary_embedding_coefficient(4 * tid + 2, rot_embed_dim, t_step, base, scale); + q.y = rotary_embedding_transform(q.y, coef1); +} + +inline __device__ void apply_rotary_embedding(bf16_4_t& q, bf16_4_t& k, int tid, int rot_embed_dim, int t_step, const float base=10000.0f, const float scale=1.0f) +{ + if (4 * tid >= rot_embed_dim) { + return; + } + const auto coef0 = rotary_embedding_coefficient(4 * tid, rot_embed_dim, t_step, base, scale); + q.x = rotary_embedding_transform(q.x, coef0); + k.x = rotary_embedding_transform(k.x, coef0); + const auto coef1 = rotary_embedding_coefficient(4 * tid + 2, rot_embed_dim, t_step, base, scale); + q.y = rotary_embedding_transform(q.y, coef1); + k.y = rotary_embedding_transform(k.y, coef1); +} + +inline __device__ void apply_rotary_embedding(bf16_8_t& q, int tid, int rot_embed_dim, int t_step, const float base=10000.0f, const float scale=1.0f) +{ + if (8 * tid >= rot_embed_dim) { + return; + } + const auto coef0 = rotary_embedding_coefficient(8 * tid, rot_embed_dim, t_step, base, scale); + q.x = rotary_embedding_transform(q.x, coef0); + const auto coef1 = rotary_embedding_coefficient(8 * tid + 2, rot_embed_dim, t_step, base, scale); + q.y = rotary_embedding_transform(q.y, coef1); + const auto coef2 = rotary_embedding_coefficient(8 * tid + 4, rot_embed_dim, t_step, base, scale); + q.z = rotary_embedding_transform(q.z, coef2); + const auto coef3 = rotary_embedding_coefficient(8 * tid + 6, rot_embed_dim, t_step, base, scale); + q.w = rotary_embedding_transform(q.w, coef3); +} + +inline __device__ void apply_rotary_embedding(bf16_8_t& q, bf16_8_t& k, int tid, int rot_embed_dim, int t_step, const float base=10000.0f, const float scale=1.0f) +{ + if (8 * tid >= rot_embed_dim) { + return; + } + const auto coef0 = rotary_embedding_coefficient(8 * tid, rot_embed_dim, t_step, base, scale); + q.x = rotary_embedding_transform(q.x, coef0); + k.x = rotary_embedding_transform(k.x, coef0); + const auto coef1 = rotary_embedding_coefficient(8 * tid + 2, rot_embed_dim, t_step, base, scale); + q.y = rotary_embedding_transform(q.y, coef1); + k.y = rotary_embedding_transform(k.y, coef1); + const auto coef2 = rotary_embedding_coefficient(8 * tid + 4, rot_embed_dim, t_step, base, scale); + q.z = rotary_embedding_transform(q.z, coef2); + k.z = rotary_embedding_transform(k.z, coef2); + const auto coef3 = rotary_embedding_coefficient(8 * tid + 6, rot_embed_dim, t_step, base, scale); + q.w = rotary_embedding_transform(q.w, coef3); + k.w = rotary_embedding_transform(k.w, coef3); +} +#endif // ENABLE_BF16 + +template +__device__ __inline__ void vec_from_smem_transpose(Vec_T& vec, T* smem, int transpose_idx, int smem_pitch); + +template<> +__device__ __inline__ void vec_from_smem_transpose(float& vec, float* smem, int transpose_idx, int smem_pitch) +{ + return; +} + +template<> +__device__ __inline__ void vec_from_smem_transpose(uint32_t& vec, uint16_t* smem, int transpose_idx, int smem_pitch) +{ + union { + uint32_t u32; + uint16_t u16[2]; + } tmp; + tmp.u16[0] = smem[transpose_idx]; + tmp.u16[1] = smem[smem_pitch + transpose_idx]; + + vec = tmp.u32; +} + +template<> +__device__ __inline__ void vec_from_smem_transpose(uint2& vec, uint16_t* smem, int transpose_idx, int smem_pitch) +{ + union { + uint32_t u32; + uint16_t u16[2]; + } tmp_1, tmp_2; + tmp_1.u32 = *reinterpret_cast(&smem[transpose_idx]); + tmp_2.u32 = *reinterpret_cast(&smem[smem_pitch + transpose_idx]); + + union { + uint2 u32x2; + uint16_t u16[4]; + } tmp_3; + tmp_3.u16[0] = tmp_1.u16[0]; + tmp_3.u16[1] = tmp_2.u16[0]; + tmp_3.u16[2] = tmp_1.u16[1]; + tmp_3.u16[3] = tmp_2.u16[1]; + + vec = tmp_3.u32x2; +} + +template<> +__device__ __inline__ void vec_from_smem_transpose(uint4& vec, uint16_t* smem, int transpose_idx, int smem_pitch) +{ + union { + uint64_t u64; + uint16_t u16[4]; + } tmp_1, tmp_2; + tmp_1.u64 = *reinterpret_cast(&smem[transpose_idx]); + tmp_2.u64 = *reinterpret_cast(&smem[smem_pitch + transpose_idx]); + + union { + uint4 u32x4; + uint16_t u16[8]; + } tmp_3; + tmp_3.u16[0] = tmp_1.u16[0]; + tmp_3.u16[1] = tmp_2.u16[0]; + tmp_3.u16[2] = tmp_1.u16[1]; + tmp_3.u16[3] = tmp_2.u16[1]; + tmp_3.u16[4] = tmp_1.u16[2]; + tmp_3.u16[5] = tmp_2.u16[2]; + tmp_3.u16[6] = tmp_1.u16[3]; + tmp_3.u16[7] = tmp_2.u16[3]; + + vec = tmp_3.u32x4; +} + +#ifdef ENABLE_BF16 +template<> +__device__ __inline__ void +vec_from_smem_transpose(bf16_4_t& vec, __nv_bfloat16* smem, int transpose_idx, int smem_pitch) +{ + union { + uint32_t u32; + __nv_bfloat16 bf16[2]; + } tmp_1, tmp_2; + tmp_1.u32 = *reinterpret_cast(&smem[transpose_idx]); + tmp_2.u32 = *reinterpret_cast(&smem[smem_pitch + transpose_idx]); + + vec.x = __nv_bfloat162{tmp_1.bf16[0], tmp_2.bf16[0]}; + vec.y = __nv_bfloat162{tmp_1.bf16[1], tmp_2.bf16[1]}; +} + +template<> +__device__ __inline__ void +vec_from_smem_transpose(bf16_8_t& vec, __nv_bfloat16* smem, int transpose_idx, int smem_pitch) +{ + union { + uint64_t u64; + __nv_bfloat16 bf16[4]; + } tmp_1, tmp_2; + tmp_1.u64 = *reinterpret_cast(&smem[transpose_idx]); + tmp_2.u64 = *reinterpret_cast(&smem[smem_pitch + transpose_idx]); + + vec.x = __nv_bfloat162{tmp_1.bf16[0], tmp_2.bf16[0]}; + vec.y = __nv_bfloat162{tmp_1.bf16[1], tmp_2.bf16[1]}; + vec.z = __nv_bfloat162{tmp_1.bf16[2], tmp_2.bf16[2]}; + vec.w = __nv_bfloat162{tmp_1.bf16[3], tmp_2.bf16[3]}; +} +#endif // ENABLE_BF16 + +template<> +__device__ __inline__ void vec_from_smem_transpose(float4& vec, float* smem, int transpose_idx, int smem_pitch) +{ + vec.x = smem[transpose_idx]; + vec.z = smem[transpose_idx + 1]; + vec.y = smem[smem_pitch + transpose_idx]; + vec.w = smem[smem_pitch + transpose_idx + 1]; +} + +template<> +__device__ __inline__ void vec_from_smem_transpose(uint32_t& vec, half* smem, int transpose_idx, int smem_pitch) +{ + union { + uint32_t u32; + half u16[2]; + } tmp; + tmp.u16[0] = smem[transpose_idx]; + tmp.u16[1] = smem[smem_pitch + transpose_idx]; + + vec = tmp.u32; +} + +#ifdef ENABLE_BF16 +template<> +__device__ __inline__ void +vec_from_smem_transpose(__nv_bfloat162& vec, __nv_bfloat16* smem, int transpose_idx, int smem_pitch) +{ + vec.x = smem[transpose_idx]; + vec.y = smem[smem_pitch + transpose_idx]; +} +#endif + +template<> +__device__ __inline__ void vec_from_smem_transpose(float2& vec, float* smem, int transpose_idx, int smem_pitch) +{ + vec.x = smem[transpose_idx]; + vec.y = smem[smem_pitch + transpose_idx]; +} + +template +__device__ __inline__ void write_smem_transpose(const Vec_T& vec, T* smem, int transpose_idx, int smem_pitch); + +template<> +__device__ __inline__ void write_smem_transpose(const float& vec, float* smem, int transpose_idx, int smem_pitch) +{ + return; +} + +template<> +__device__ __inline__ void write_smem_transpose(const uint4& vec, uint16_t* smem, int transpose_idx, int smem_pitch) +{ + union { + uint64_t u64; + uint16_t u16[4]; + } tmp_1, tmp_2; + + union { + uint4 u32x4; + uint16_t u16[8]; + } tmp_3; + tmp_3.u32x4 = vec; + tmp_1.u16[0] = tmp_3.u16[0]; + tmp_2.u16[0] = tmp_3.u16[1]; + tmp_1.u16[1] = tmp_3.u16[2]; + tmp_2.u16[1] = tmp_3.u16[3]; + tmp_1.u16[2] = tmp_3.u16[4]; + tmp_2.u16[2] = tmp_3.u16[5]; + tmp_1.u16[3] = tmp_3.u16[6]; + tmp_2.u16[3] = tmp_3.u16[7]; + + *reinterpret_cast(&smem[transpose_idx]) = tmp_1.u64; + *reinterpret_cast(&smem[smem_pitch + transpose_idx]) = tmp_2.u64; +} + +template<> +__device__ __inline__ void write_smem_transpose(const uint2& vec, uint16_t* smem, int transpose_idx, int smem_pitch) +{ + union { + uint32_t u32; + uint16_t u16[2]; + } tmp_1, tmp_2; + + union { + uint2 u32x2; + uint16_t u16[4]; + } tmp_3; + tmp_3.u32x2 = vec; + tmp_1.u16[0] = tmp_3.u16[0]; + tmp_2.u16[0] = tmp_3.u16[1]; + tmp_1.u16[1] = tmp_3.u16[2]; + tmp_2.u16[1] = tmp_3.u16[3]; + + *reinterpret_cast(&smem[transpose_idx]) = tmp_1.u32; + *reinterpret_cast(&smem[smem_pitch + transpose_idx]) = tmp_2.u32; +} + +template<> +__device__ __inline__ void write_smem_transpose(const uint32_t& vec, uint16_t* smem, int transpose_idx, int smem_pitch) +{ + union { + uint32_t u32; + uint16_t u16[2]; + } tmp; + tmp.u32 = vec; + + smem[transpose_idx] = tmp.u16[0]; + smem[smem_pitch + transpose_idx] = tmp.u16[1]; +} + +template<> +__device__ __inline__ void write_smem_transpose(const float4& vec, float* smem, int transpose_idx, int smem_pitch) +{ + smem[transpose_idx] = vec.x; + smem[transpose_idx + 1] = vec.z; + smem[smem_pitch + transpose_idx] = vec.y; + smem[smem_pitch + transpose_idx + 1] = vec.w; +} + +template<> +__device__ __inline__ void write_smem_transpose(const uint32_t& vec, half* smem, int transpose_idx, int smem_pitch) +{ + union { + uint32_t u32; + half u16[2]; + } tmp; + + tmp.u32 = vec; + smem[transpose_idx] = tmp.u16[0]; + smem[smem_pitch + transpose_idx] = tmp.u16[1]; +} + +#ifdef ENABLE_BF16 +template<> +__device__ __inline__ void +write_smem_transpose(const __nv_bfloat162& vec, __nv_bfloat16* smem, int transpose_idx, int smem_pitch) +{ + smem[transpose_idx] = vec.x; + smem[smem_pitch + transpose_idx] = vec.y; +} + +template<> +__device__ __inline__ void +write_smem_transpose(const bf16_4_t& vec, __nv_bfloat16* smem, int transpose_idx, int smem_pitch) +{ + write_smem_transpose(reinterpret_cast(vec), reinterpret_cast(smem), transpose_idx, smem_pitch); +} + +template<> +__device__ __inline__ void +write_smem_transpose(const bf16_8_t& vec, __nv_bfloat16* smem, int transpose_idx, int smem_pitch) +{ + write_smem_transpose(reinterpret_cast(vec), reinterpret_cast(smem), transpose_idx, smem_pitch); +} +#endif + +template<> +__device__ __inline__ void write_smem_transpose(const float2& vec, float* smem, int transpose_idx, int smem_pitch) +{ + smem[transpose_idx] = vec.x; + smem[smem_pitch + transpose_idx] = vec.y; +} + +} // namespace mmha diff --git a/llm-awq/awq/kernels/csrc/attention/ft_attention.cpp b/llm-awq/awq/kernels/csrc/attention/ft_attention.cpp new file mode 100644 index 0000000000000000000000000000000000000000..37d25bad49963ea1797cdd2b58202aefe1bcec13 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/attention/ft_attention.cpp @@ -0,0 +1,185 @@ +// Adapted from NVIDIA/FasterTransformer and FlashAttention + +#include +#include "ATen/cuda/CUDAContext.h" +#include + +#include "ft_attention.h" +#include "decoder_masked_multihead_attention.h" + +#define CHECK_DEVICE(x) TORCH_CHECK(x.device().type() == torch::kCUDA, #x " must be on CUDA") +#define CHECK_SHAPE(x, ...) TORCH_CHECK(x.sizes() == torch::IntArrayRef({__VA_ARGS__}), #x " must have shape (" #__VA_ARGS__ ")") +#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous") + +#define DISPATCH_FLOAT_AND_HALF_AND_BF16(TYPE, NAME, ...) \ + if (TYPE == at::ScalarType::Half) { \ + using scalar_t = at::Half; \ + __VA_ARGS__(); \ + } else if (TYPE == at::ScalarType::BFloat16) { \ + using scalar_t = at::BFloat16; \ + __VA_ARGS__(); \ + } else if (TYPE == at::ScalarType::Float) { \ + using scalar_t = float; \ + __VA_ARGS__(); \ + } else { \ + AT_ERROR(#NAME, " not implemented for type '", toString(TYPE), "'"); \ + } + +template +void masked_multihead_attention(const Masked_multihead_attention_params& params, + const cudaStream_t& stream); + +template +void cross_multihead_attention(const Masked_multihead_attention_params& params, + const cudaStream_t& stream); + +template +struct SATypeConverter { + using Type = T; +}; + +template<> +struct SATypeConverter { + using Type = uint16_t; +}; + +template<> +struct SATypeConverter { + using Type = __nv_bfloat16; +}; + +template +void set_params(Masked_multihead_attention_params ¶ms, + const size_t batch_size, + const size_t nheads, + const size_t nheads_kv, + const size_t memory_max_seqlen, + const size_t headdim, + const int timestep, + const int rotary_embedding_dim, + const float rotary_base, + const float rotary_scale, + const bool neox_rotary_style, + const int qkv_batch_stride, + T *q_ptr, + T *k_ptr, + T *v_ptr, + T *k_cache_ptr, + T *v_cache_ptr, + int *length_per_sample, + float *alibi_slopes_ptr, + T *out_ptr) { + // Reset the parameters + memset(¶ms, 0, sizeof(params)); + params.q = q_ptr; + params.k = k_ptr; + params.v = v_ptr; + params.q_bias = nullptr; + params.k_bias = nullptr; + params.v_bias = nullptr; + params.k_cache = k_cache_ptr; + params.v_cache = v_cache_ptr; + params.linear_bias_slopes = alibi_slopes_ptr; + params.out = out_ptr; + params.cache_indir = nullptr; + params.stride = qkv_batch_stride; + params.batch_size = batch_size; + params.beam_width = 1; + params.memory_max_len = memory_max_seqlen; + params.num_heads = nheads; + params.num_kv_heads = nheads_kv; + params.hidden_size_per_head = headdim; + params.rotary_embedding_dim = rotary_embedding_dim; + params.rotary_base = rotary_base; + params.rotary_scale = rotary_scale; + params.neox_rotary_style = neox_rotary_style; + params.timestep = timestep; + params.inv_sqrt_dh = 1.f / sqrt(float(headdim)); + params.total_padding_tokens = nullptr; + params.masked_tokens = nullptr; + params.prefix_prompt_lengths = nullptr; + params.max_prefix_prompt_length = 0; + params.relative_attention_bias = nullptr; + params.relative_attention_bias_stride = 0; + params.cross_attention_out = nullptr; + params.max_decoder_seq_len = 0; + params.is_return_cross_attentions = false; + params.finished = nullptr; + params.memory_length_per_sample = nullptr; + params.length_per_sample = length_per_sample; +} + +torch::Tensor single_query_attention(const torch::Tensor q, + const torch::Tensor k, + const torch::Tensor v, + torch::Tensor k_cache, + torch::Tensor v_cache, + c10::optional length_per_sample_, + c10::optional alibi_slopes_, + const int timestep, + const int rotary_embedding_dim, + const float rotary_base, + const float rotary_scale, + // neox_rotary_style = not interleaved + const bool neox_rotary_style) { + CHECK_DEVICE(q); CHECK_DEVICE(k); CHECK_DEVICE(v); CHECK_DEVICE(k_cache); CHECK_DEVICE(v_cache); + int batch_size = v_cache.size(0); + int nheads = q.size(1); + int nheads_kv = v_cache.size(1); + int memory_max_seqlen = v_cache.size(2); + int headdim = v_cache.size(3); + CHECK_SHAPE(q, batch_size, nheads, headdim); + CHECK_SHAPE(k, batch_size, nheads_kv, headdim); + CHECK_SHAPE(v, batch_size, nheads_kv, headdim); + CHECK_SHAPE(v_cache, batch_size, nheads_kv, memory_max_seqlen, headdim); + // k_cache shape: [B, H, Dh/x, L, x] where x=8 for fp16 and x=4 for fp32 + int packsize = k_cache.dtype() == torch::kFloat32 ? 4 : 8; + CHECK_SHAPE(k_cache, batch_size, nheads_kv, headdim / packsize, memory_max_seqlen, packsize); + TORCH_CHECK(q.stride(2) == 1 && q.stride(1) == headdim); + TORCH_CHECK(k.stride(2) == 1 && k.stride(1) == headdim); + TORCH_CHECK(v.stride(2) == 1 && v.stride(1) == headdim); + // TORCH_CHECK(q.stride(0) == k.stride(0) && q.stride(0) == v.stride(0)); + CHECK_CONTIGUOUS(v_cache); CHECK_CONTIGUOUS(k_cache); + + if (length_per_sample_.has_value()) { + auto length_per_sample = length_per_sample_.value(); + CHECK_DEVICE(length_per_sample); + CHECK_SHAPE(length_per_sample, batch_size); + CHECK_CONTIGUOUS(length_per_sample); + TORCH_CHECK(length_per_sample.dtype() == torch::kInt32); + } + + if (alibi_slopes_.has_value()) { + auto alibi_slopes = alibi_slopes_.value(); + CHECK_DEVICE(alibi_slopes); + CHECK_SHAPE(alibi_slopes, nheads); + CHECK_CONTIGUOUS(alibi_slopes); + TORCH_CHECK(alibi_slopes.dtype() == torch::kFloat32); + } + + // Otherwise the kernel will be launched from cuda:0 device + // Cast to char to avoid compiler warning about narrowing + at::cuda::CUDAGuard device_guard{(char)q.get_device()}; + + torch::Tensor out = torch::empty_like(q); + + DISPATCH_FLOAT_AND_HALF_AND_BF16(q.scalar_type(), "single_query_attention", [&] { + using DataType = typename SATypeConverter::Type; + Masked_multihead_attention_params params; + set_params(params, batch_size, nheads, nheads_kv, memory_max_seqlen, headdim, + timestep, rotary_embedding_dim, rotary_base, rotary_scale, neox_rotary_style, q.stride(0), + reinterpret_cast(q.data_ptr()), + reinterpret_cast(k.data_ptr()), + reinterpret_cast(v.data_ptr()), + reinterpret_cast(k_cache.data_ptr()), + reinterpret_cast(v_cache.data_ptr()), + length_per_sample_.has_value() + ? length_per_sample_.value().data_ptr() : nullptr, + alibi_slopes_.has_value() + ? alibi_slopes_.value().data_ptr(): nullptr, + reinterpret_cast(out.data_ptr())); + auto stream = at::cuda::getCurrentCUDAStream(); + masked_multihead_attention(params, stream); + }); + return out; +} \ No newline at end of file diff --git a/llm-awq/awq/kernels/csrc/attention/ft_attention.h b/llm-awq/awq/kernels/csrc/attention/ft_attention.h new file mode 100644 index 0000000000000000000000000000000000000000..53037116aae1d7858c63870125a146b4c2f7df87 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/attention/ft_attention.h @@ -0,0 +1,16 @@ +#pragma once +#include + + +torch::Tensor single_query_attention(const torch::Tensor q, + const torch::Tensor k, + const torch::Tensor v, + torch::Tensor k_cache, + torch::Tensor v_cache, + c10::optional length_per_sample_, + c10::optional alibi_slopes_, + const int timestep, + const int rotary_embedding_dim = 0, + const float rotary_base = 10000.0f, + const float rotary_scale = 1.0f, + const bool neox_rotary_style=true); \ No newline at end of file diff --git a/llm-awq/awq/kernels/csrc/attention/setup.py b/llm-awq/awq/kernels/csrc/attention/setup.py new file mode 100644 index 0000000000000000000000000000000000000000..6c19b38b9cebf7a44634e31442ba93639cfbd72b --- /dev/null +++ b/llm-awq/awq/kernels/csrc/attention/setup.py @@ -0,0 +1,159 @@ +# Adapted from https://github.com/NVIDIA/apex/blob/master/setup.py +import sys +import warnings +import os +from packaging.version import parse, Version + +from setuptools import setup, find_packages +import subprocess + +import torch +from torch.utils.cpp_extension import ( + BuildExtension, + CppExtension, + CUDAExtension, + CUDA_HOME, +) + + +# ninja build does not work unless include_dirs are abs path +this_dir = os.path.dirname(os.path.abspath(__file__)) + + +def get_cuda_bare_metal_version(cuda_dir): + raw_output = subprocess.check_output( + [cuda_dir + "/bin/nvcc", "-V"], universal_newlines=True + ) + output = raw_output.split() + release_idx = output.index("release") + 1 + bare_metal_version = parse(output[release_idx].split(",")[0]) + + return raw_output, bare_metal_version + + +def check_cuda_torch_binary_vs_bare_metal(cuda_dir): + raw_output, bare_metal_version = get_cuda_bare_metal_version(cuda_dir) + torch_binary_version = parse(torch.version.cuda) + + print("\nCompiling cuda extensions with") + print(raw_output + "from " + cuda_dir + "/bin\n") + + if bare_metal_version != torch_binary_version: + raise RuntimeError( + "Cuda extensions are being compiled with a version of Cuda that does " + "not match the version used to compile Pytorch binaries. " + "Pytorch binaries were compiled with Cuda {}.\n".format(torch.version.cuda) + + "In some cases, a minor-version mismatch will not cause later errors: " + "https://github.com/NVIDIA/apex/pull/323#discussion_r287021798. " + "You can try commenting out this check (at your own risk)." + ) + + +def raise_if_cuda_home_none(global_option: str) -> None: + if CUDA_HOME is not None: + return + raise RuntimeError( + f"{global_option} was requested, but nvcc was not found. Are you sure your environment has nvcc available? " + "If you're installing within a container from https://hub.docker.com/r/pytorch/pytorch, " + "only images whose names contain 'devel' will provide nvcc." + ) + + +def append_nvcc_threads(nvcc_extra_args): + _, bare_metal_version = get_cuda_bare_metal_version(CUDA_HOME) + if bare_metal_version >= Version("11.2"): + return nvcc_extra_args + ["--threads", "4"] + return nvcc_extra_args + + +if not torch.cuda.is_available(): + # https://github.com/NVIDIA/apex/issues/486 + # Extension builds after https://github.com/pytorch/pytorch/pull/23408 attempt to query torch.cuda.get_device_capability(), + # which will fail if you are compiling in an environment without visible GPUs (e.g. during an nvidia-docker build command). + print( + "\nWarning: Torch did not find available GPUs on this system.\n", + "If your intention is to cross-compile, this is not an error.\n" + "By default, Apex will cross-compile for Pascal (compute capabilities 6.0, 6.1, 6.2),\n" + "Volta (compute capability 7.0), Turing (compute capability 7.5),\n" + "and, if the CUDA version is >= 11.0, Ampere (compute capability 8.0).\n" + "If you wish to cross-compile for a single specific architecture,\n" + 'export TORCH_CUDA_ARCH_LIST="compute capability" before running setup.py.\n', + ) + if os.environ.get("TORCH_CUDA_ARCH_LIST", None) is None and CUDA_HOME is not None: + _, bare_metal_version = get_cuda_bare_metal_version(CUDA_HOME) + if bare_metal_version >= Version("11.8"): + os.environ["TORCH_CUDA_ARCH_LIST"] = "6.0;6.1;6.2;7.0;7.5;8.0;8.6;9.0" + elif bare_metal_version >= Version("11.1"): + os.environ["TORCH_CUDA_ARCH_LIST"] = "6.0;6.1;6.2;7.0;7.5;8.0;8.6" + elif bare_metal_version == Version("11.0"): + os.environ["TORCH_CUDA_ARCH_LIST"] = "6.0;6.1;6.2;7.0;7.5;8.0" + else: + os.environ["TORCH_CUDA_ARCH_LIST"] = "6.0;6.1;6.2;7.0;7.5" + + +print("\n\ntorch.__version__ = {}\n\n".format(torch.__version__)) +TORCH_MAJOR = int(torch.__version__.split(".")[0]) +TORCH_MINOR = int(torch.__version__.split(".")[1]) + +cmdclass = {} +ext_modules = [] + +# Check, if ATen/CUDAGeneratorImpl.h is found, otherwise use ATen/cuda/CUDAGeneratorImpl.h +# See https://github.com/pytorch/pytorch/pull/70650 +generator_flag = [] +torch_dir = torch.__path__[0] +if os.path.exists(os.path.join(torch_dir, "include", "ATen", "CUDAGeneratorImpl.h")): + generator_flag = ["-DOLD_GENERATOR_PATH"] + +raise_if_cuda_home_none("--ft_attention") +# Check, if CUDA11 is installed for compute capability 8.0 +cc_flag = [] +_, bare_metal_version = get_cuda_bare_metal_version(CUDA_HOME) +if bare_metal_version < Version("11.0"): + raise RuntimeError("ft_attention is only supported on CUDA 11 and above") +cc_flag.append("-gencode") +cc_flag.append("arch=compute_70,code=sm_70") +cc_flag.append("-gencode") +cc_flag.append("arch=compute_80,code=sm_80") +if bare_metal_version >= Version("11.8"): + cc_flag.append("-gencode") + cc_flag.append("arch=compute_90,code=sm_90") + +ext_modules.append( + CUDAExtension( + name="ft_attention", + sources=[ + "ft_attention.cpp", + "decoder_masked_multihead_attention.cu", + ], + extra_compile_args={ + "cxx": ["-O3", "-DENABLE_BF16"] + generator_flag, + "nvcc": append_nvcc_threads( + [ + "-DENABLE_BF16", # TODO + "-O3", + "-U__CUDA_NO_HALF_OPERATORS__", + "-U__CUDA_NO_HALF_CONVERSIONS__", + "-U__CUDA_NO_BFLOAT16_OPERATORS__", + "-U__CUDA_NO_BFLOAT16_CONVERSIONS__", + "-U__CUDA_NO_BFLOAT162_OPERATORS__", + "-U__CUDA_NO_BFLOAT162_CONVERSIONS__", + "--expt-relaxed-constexpr", + "--expt-extended-lambda", + "--use_fast_math", + ] + + generator_flag + + cc_flag + ), + }, + include_dirs=[this_dir], + ) +) + +setup( + name="ft_attention", + version="0.1", + description="Attention for single query from FasterTransformer", + ext_modules=ext_modules, + cmdclass={"build_ext": BuildExtension} if ext_modules else {}, +) diff --git a/llm-awq/awq/kernels/csrc/layernorm/layernorm.cu b/llm-awq/awq/kernels/csrc/layernorm/layernorm.cu new file mode 100644 index 0000000000000000000000000000000000000000..8f2de9a199f5f7f2bcbeac050da7dfce01192eba --- /dev/null +++ b/llm-awq/awq/kernels/csrc/layernorm/layernorm.cu @@ -0,0 +1,131 @@ +/* + +Adapted from NVIDIA FasterTransformer: +https://github.com/NVIDIA/FasterTransformer/blob/main/src/fastertransformer/kernels/layernorm_kernels.cu + +*/ + +#include +#include +#include "reduction.cuh" +#include "layernorm.h" +#include +#include + +#define DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP16(pytorch_dtype, c_type, ...) \ + if (pytorch_dtype == at::ScalarType::Half) { \ + using c_type = half; \ + __VA_ARGS__ \ + } else if (pytorch_dtype == at::ScalarType::BFloat16) { \ + using c_type = nv_bfloat16; \ + __VA_ARGS__ \ + } else { \ + std::ostringstream oss; \ + oss << __PRETTY_FUNCTION__ << " failed to dispatch data type " << pytorch_dtype; \ + TORCH_CHECK(false, oss.str()); \ + } + +static inline __device__ float to_float(half src) +{ + return __half2float(src); +} + +static inline __device__ float to_float(float src) +{ + return src; +} + +template +__global__ void generalT5LayerNorm( + const T* __restrict input, const T* __restrict gamma, T* output, const float layernorm_eps, int m, int n) +{ + // layernorm module in the T5 style No bias and no subtraction of mean. + const int tid = threadIdx.x; + + __shared__ float s_variance; + float variance = 0.0f; + + float local_var_sum = 0.0f; + for (int i = tid; i < n; i += blockDim.x) { + float diff = to_float(__ldg(&input[blockIdx.x * n + i])); + local_var_sum += diff * diff; + } + variance = blockReduceSum(local_var_sum); + + if (threadIdx.x == 0) { + s_variance = rsqrtf(variance / (float)n + layernorm_eps); + } + __syncthreads(); + + for (int i = tid; i < n; i += blockDim.x) { + output[blockIdx.x * n + i] = + clamp_inf_for_half((to_float(input[blockIdx.x * n + i]) * s_variance) * to_float(__ldg(&gamma[i]))); + } +} + + +template +void invokeGeneralT5LayerNorm(T* out, + const T* input, + const T* gamma, + // const T* beta, + const float layernorm_eps, + const int m, + const int n) +{ + dim3 grid(m); + dim3 block(min(n, 1024)); + + /* For general cases, n is equal to hidden_units, e.g., 512/1024. + Since we have warp shuffle inside the code, block.x % 32 should be 0. + */ + if (n % 32 != 0) { + block.x = 1024; + } + + block.x = block.x / (4 / sizeof(T)); // if using half, only need half of block.x + + /* should pay attention to the rsqrt precision*/ + generalT5LayerNorm<<>>(input, gamma, out, layernorm_eps, m, n); // For gpt-3 +} + +template void invokeGeneralT5LayerNorm(half* out, + const half* input, + const half* gamma, + // const half* beta, + const float layernorm_eps, + const int m, + const int n); + +template void invokeGeneralT5LayerNorm(float* out, + const float* input, + const float* gamma, + // const half* beta, + const float layernorm_eps, + const int m, + const int n); + + + +// input b, n, c +void layernorm_forward_cuda( + torch::Tensor _input, + torch::Tensor _gamma, + torch::Tensor _out, + float eps) +{ + int m = _input.size(0) * _input.size(1); + int n = _input.size(2); + const at::cuda::OptionalCUDAGuard device_guard(device_of(_input)); + + auto data_type = _input.scalar_type(); + TORCH_CHECK(_gamma.scalar_type() == data_type); + TORCH_CHECK(_out.scalar_type() == data_type); + + DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP16(data_type, ctype, { + auto input = reinterpret_cast(_input.data_ptr()); + auto gamma = reinterpret_cast(_gamma.data_ptr()); + auto out = reinterpret_cast(_out.data_ptr()); + invokeGeneralT5LayerNorm(out, input, gamma, eps, m, n); + }); +} diff --git a/llm-awq/awq/kernels/csrc/layernorm/layernorm.h b/llm-awq/awq/kernels/csrc/layernorm/layernorm.h new file mode 100644 index 0000000000000000000000000000000000000000..de43ccac688d65b540b8fc9838e9e3c44b1758a1 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/layernorm/layernorm.h @@ -0,0 +1,3 @@ +#include + +void layernorm_forward_cuda(torch::Tensor _input, torch::Tensor _gamma, torch::Tensor _out, float eps); diff --git a/llm-awq/awq/kernels/csrc/layernorm/reduction.cuh b/llm-awq/awq/kernels/csrc/layernorm/reduction.cuh new file mode 100644 index 0000000000000000000000000000000000000000..678160e8fdf5788757a82060bab1ca6b6f6d3baf --- /dev/null +++ b/llm-awq/awq/kernels/csrc/layernorm/reduction.cuh @@ -0,0 +1,82 @@ +/* + +Adapted from NVIDIA FasterTransformer: +https://github.com/NVIDIA/FasterTransformer/blob/main/src/fastertransformer/kernels/reduce_kernel_utils.cuh +*/ + +#pragma once +#include +#if ((__CUDACC_VER_MAJOR__ > 11) || (__CUDACC_VER_MAJOR__ == 11 && __CUDACC_VER_MINOR__ >= 0)) +#include +#else +#include +#endif +#include +#include +#include +#include + +static const float HALF_FLT_MAX = 65504.F; +#define FINAL_MASK 0xffffffff + + +template +inline __device__ T add(T a, T b) { + return a + b; +} + +template<> +inline __device__ half2 add(half2 a, half2 b) { + return __hadd2(a, b); +} + +template<> +inline __device__ half add(half a, half b) { + return __hadd(a, b); +} + +template +__inline__ __device__ T warpReduceSum(T val) +{ +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) + val = add(val, __shfl_xor_sync(FINAL_MASK, val, mask, 32)); //__shfl_sync bf16 return float when sm < 80 + return val; +} + +/* Calculate the sum of all elements in a block */ +template +__inline__ __device__ T blockReduceSum(T val) +{ + static __shared__ T shared[32]; + int lane = threadIdx.x & 0x1f; + int wid = threadIdx.x >> 5; + + val = warpReduceSum(val); + + if (lane == 0) + shared[wid] = val; + + __syncthreads(); + + // Modify from blockDim.x << 5 to blockDim.x / 32. to prevent + // blockDim.x is not divided by 32 + val = (threadIdx.x < (blockDim.x / 32.f)) ? shared[lane] : (T)(0.0f); + val = warpReduceSum(val); + + return val; +} + + +template +__device__ __forceinline__ T clamp_inf_for_half(const float input) +{ + return input; +} + +template<> +__device__ __forceinline__ half clamp_inf_for_half(const float input) +{ + // clamp inf values to enable fp16 training + return input > 0.0f ? __float2half(min(input, HALF_FLT_MAX - 1000)) : __float2half(max(input, -HALF_FLT_MAX + 1000)); +} diff --git a/llm-awq/awq/kernels/csrc/position_embedding/pos_encoding.h b/llm-awq/awq/kernels/csrc/position_embedding/pos_encoding.h new file mode 100644 index 0000000000000000000000000000000000000000..04e205b238c56142a5568294978338c592a2005d --- /dev/null +++ b/llm-awq/awq/kernels/csrc/position_embedding/pos_encoding.h @@ -0,0 +1,9 @@ +#pragma once +#include + +void rotary_embedding_neox( + torch::Tensor& positions, + torch::Tensor& query, + torch::Tensor& key, + int head_size, + torch::Tensor& cos_sin_cache); \ No newline at end of file diff --git a/llm-awq/awq/kernels/csrc/position_embedding/pos_encoding_kernels.cu b/llm-awq/awq/kernels/csrc/position_embedding/pos_encoding_kernels.cu new file mode 100644 index 0000000000000000000000000000000000000000..883b59c41b74cf1947c1b1e1bacde716f0f21857 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/position_embedding/pos_encoding_kernels.cu @@ -0,0 +1,88 @@ +/* + +Adapted from the VLLM project: +https://github.com/vllm-project/vllm/blob/main/csrc/pos_encoding_kernels.cu + +*/ + +#include +#include +#include "pos_encoding.h" + +template +__global__ void rotary_embedding_neox_kernel( + const int64_t* __restrict__ positions, // [num_tokens] + scalar_t* __restrict__ query, // [num_tokens, num_heads, head_size] + scalar_t* __restrict__ key, // [num_tokens, num_heads, head_size] + const scalar_t* __restrict__ cos_sin_cache, // [max_position, 2, rot_dim // 2] + const int rot_dim, + const int stride, + const int num_heads, + const int head_size) { + // Each thread block is responsible for one token. + const int token_idx = blockIdx.x; + int64_t pos = positions[token_idx]; + const scalar_t* cache_ptr = cos_sin_cache + pos * rot_dim; + + const int embed_dim = rot_dim / 2; + const int n = num_heads * embed_dim; + for (int i = threadIdx.x; i < n; i += blockDim.x) { + const int head_idx = i / embed_dim; + const int token_head = token_idx * stride + head_idx * head_size; + + const int rot_offset = i % embed_dim; + const int x_index = rot_offset; + const int y_index = embed_dim + rot_offset; + + const int out_x = token_idx * stride + head_idx * head_size + x_index; + const int out_y = token_idx * stride + head_idx * head_size + y_index; + + const scalar_t cos = __ldg(cache_ptr + x_index); + const scalar_t sin = __ldg(cache_ptr + y_index); + + const scalar_t q_x = query[token_head + x_index]; + const scalar_t q_y = query[token_head + y_index]; + query[out_x] = q_x * cos - q_y * sin; + query[out_y] = q_y * cos + q_x * sin; + + const scalar_t k_x = key[token_head + x_index]; + const scalar_t k_y = key[token_head + y_index]; + key[out_x] = k_x * cos - k_y * sin; + key[out_y] = k_y * cos + k_x * sin; + } +} + +void rotary_embedding_neox( + torch::Tensor& positions, // [b, num_tokens] + torch::Tensor& query, // [b, num_tokens, 1, num_heads, head_size] + torch::Tensor& key, // [b, num_tokens, 1, num_heads, head_size] + int head_size, + torch::Tensor& cos_sin_cache) // [max_position, rot_dim] +{ + int num_tokens = query.size(0) * query.size(1); + int rot_dim = cos_sin_cache.size(1); + int num_heads = query.size(-2); + int stride = num_heads * head_size; + // TORCH_CHECK(stride == key.stride(0)); + + dim3 grid(num_tokens); + dim3 block(std::min(num_heads * rot_dim / 2, 512)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, + at::ScalarType::BFloat16, + query.scalar_type(), + "rotary_embedding_neox", + [&] { + rotary_embedding_neox_kernel<<>>( + positions.data_ptr(), + query.data_ptr(), + key.data_ptr(), + cos_sin_cache.data_ptr(), + rot_dim, + stride, + num_heads, + head_size); + }); +} + diff --git a/llm-awq/awq/kernels/csrc/pybind.cpp b/llm-awq/awq/kernels/csrc/pybind.cpp new file mode 100644 index 0000000000000000000000000000000000000000..30424ee05147bcb5d66f2df2a7819470934b7051 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/pybind.cpp @@ -0,0 +1,38 @@ +#include +#include +#include "attention/ft_attention.h" +#include "layernorm/layernorm.h" +#include "quantization/gemm_cuda.h" +#include "quantization/gemv_cuda.h" +#include "quantization_new/gemm/gemm_cuda.h" +#include "quantization_new/gemv/gemv_cuda.h" +#include "position_embedding/pos_encoding.h" +#include "rope_new/fused_rope_with_pos.h" +#include "w8a8/w8a8_gemm_cuda.h" +#include "w8a8/quantization.h" +#include "w8a8/layernorm.h" +#include "w8a8/act.h" + + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + m.def("layernorm_forward_cuda", &layernorm_forward_cuda, "FasterTransformer layernorm kernel"); + m.def("gemm_forward_cuda", &gemm_forward_cuda, "Quantized GEMM kernel."); + m.def("gemv_forward_cuda", &gemv_forward_cuda, "Quantized GEMV kernel."); + m.def("gemm_forward_cuda_new", &gemm_forward_cuda_new, "New quantized GEMM kernel."); + m.def("gemv_forward_cuda_new", &gemv_forward_cuda_new, "New quantized GEMV kernel."); + m.def("rotary_embedding_neox", &rotary_embedding_neox, "Apply GPT-NeoX style rotary embedding to query and key"); + m.def("single_query_attention", &single_query_attention, "Attention with a single query", + py::arg("q"), py::arg("k"), py::arg("v"), py::arg("k_cache"), py::arg("v_cache"), + py::arg("length_per_sample_"), py::arg("alibi_slopes_"), py::arg("timestep"), py::arg("rotary_embedding_dim")=0, + py::arg("rotary_base")=10000.0f, py::arg("rotary_scale")=1.0f, py::arg("neox_rotary_style")=true); + m.def("fused_rope_with_pos_forward_func", &fused_rope_with_pos_forward_func,"Fused rope forward function with B,S,D embedding"); + m.def("w8a8_gemm_forward_cuda", &w8a8_gemm_forward_cuda, "our w8a8 gemm kernel"); + m.def("w8a8_gemm_fuse_bias_forward_cuda", &w8a8_gemm_fuse_bias_forward_cuda, "our w8a8 gemm fused bias kernel"); + m.def("invoke_quant", &invoke_quant, "fp16->int8 quantization"); + m.def("rms_norm_general", &rms_norm_general, py::arg("out"), py::arg("input"), + py::arg("weight"), py::arg("bias"),py::arg("scaling"), py::arg("epsilon"), py::arg("use_per_token_quant") = true, + "Apply Root Mean Square (RMS) Normalization to the input tensor (TRTLLM kernel)."); + m.def("silu_and_mul", &silu_and_mul, "Activation function."); + m.def("gelu_and_quant",&gelu_and_quant, "Apply gelu act and quant output"); +} diff --git a/llm-awq/awq/kernels/csrc/quantization/dequantize.cuh b/llm-awq/awq/kernels/csrc/quantization/dequantize.cuh new file mode 100644 index 0000000000000000000000000000000000000000..5d333b35c148d0cd01c8cb6fffd5deb28db28d33 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/quantization/dequantize.cuh @@ -0,0 +1,79 @@ +/* +Modified from NVIDIA FasterTransformer: https://github.com/NVIDIA/FasterTransformer/blob/main/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h + +@article{lin2023awq, + title={AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration}, + author={Lin, Ji and Tang, Jiaming and Tang, Haotian and Yang, Shang and Dang, Xingyu and Han, Song}, + journal={arXiv}, + year={2023} +} +*/ + +#pragma once + + +__device__ uint4 dequantize_s4_to_fp16x2(uint32_t const& source) +{ + uint4 result; + + uint32_t* h = reinterpret_cast(&result); + uint32_t const i4s = reinterpret_cast(source); + + // First, we extract the i4s and construct an intermediate fp16 number. + static constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa; + static constexpr uint32_t BOTTOM_MASK = 0x000f000f; + static constexpr uint32_t TOP_MASK = 0x00f000f0; + static constexpr uint32_t I4s_TO_F16s_MAGIC_NUM = 0x64006400; + + // Note that the entire sequence only requires 1 shift instruction. This is thanks to the register packing + // format and the fact that we force our integers to be unsigned, and account for this in the fp16 subtractions. + // In addition, I exploit the fact that sub and fma have the same throughput in order to convert elt_23 and + // elt_67 to fp16 without having to shift them to the bottom bits before hand. + + // Shift right by 8 to now consider elt_45 and elt_67. Issue first to hide RAW dependency if we issue + // immediately before required. + const uint32_t top_i4s = i4s >> 8; + // Extract elt_01 - (i4s & 0x000f000f) | 0x64006400 + asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(h[0]) + : "r"(i4s), "n"(BOTTOM_MASK), "n"(I4s_TO_F16s_MAGIC_NUM), "n"(immLut)); + // Extract elt_23 (i4s & 0x00f000f0) | 0x64006400 + asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(h[1]) + : "r"(i4s), "n"(TOP_MASK), "n"(I4s_TO_F16s_MAGIC_NUM), "n"(immLut)); + // Extract elt_45 (top_i4s & 0x000f000f) | 0x64006400 + asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(h[2]) + : "r"(top_i4s), "n"(BOTTOM_MASK), "n"(I4s_TO_F16s_MAGIC_NUM), "n"(immLut)); + // Extract elt_67 (top_i4s & 0x00f000f0) | 0x64006400 + asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(h[3]) + : "r"(top_i4s), "n"(TOP_MASK), "n"(I4s_TO_F16s_MAGIC_NUM), "n"(immLut)); + + // I use inline PTX below because I am not sure if the compiler will emit float2half instructions if I use the + // half2 ctor. In this case, I chose performance reliability over code readability. + + // This is the half2 {1032, 1032} represented as an integer. + // static constexpr uint32_t FP16_TOP_MAGIC_NUM = 0x64086408; + // Haotian: subtract {1024, 1024} instead, we do not need to map to [-8, 7] + static constexpr uint32_t FP16_TOP_MAGIC_NUM = 0x64006400; + // This is the half2 {1 / 16, 1 / 16} represented as an integer. + static constexpr uint32_t ONE_SIXTEENTH = 0x2c002c00; + // This is the half2 {-72, -72} represented as an integer. + // static constexpr uint32_t NEG_72 = 0xd480d480; + // Haotian: Let's use {-64, -64}. + static constexpr uint32_t NEG_64 = 0xd400d400; + + // Finally, we construct the output numbers. + // Convert elt_01 + asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(h[0]) : "r"(h[0]), "r"(FP16_TOP_MAGIC_NUM)); + // Convert elt_23 + asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n" : "=r"(h[1]) : "r"(h[1]), "r"(ONE_SIXTEENTH), "r"(NEG_64)); + // Convert elt_45 + asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(h[2]) : "r"(h[2]), "r"(FP16_TOP_MAGIC_NUM)); + // Convert elt_67 + asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n" : "=r"(h[3]) : "r"(h[3]), "r"(ONE_SIXTEENTH), "r"(NEG_64)); + + return result; +} + diff --git a/llm-awq/awq/kernels/csrc/quantization/gemm_cuda.h b/llm-awq/awq/kernels/csrc/quantization/gemm_cuda.h new file mode 100644 index 0000000000000000000000000000000000000000..4c0846c84866b23a44dc9a640358bf54b160bec9 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/quantization/gemm_cuda.h @@ -0,0 +1,4 @@ +#include + +torch::Tensor gemm_forward_cuda(torch::Tensor _in_feats, torch::Tensor _kernel, + torch::Tensor _scaling_factors, torch::Tensor _zeros, int group_size, int split_k_iters); diff --git a/llm-awq/awq/kernels/csrc/quantization/gemm_cuda_gen.cu b/llm-awq/awq/kernels/csrc/quantization/gemm_cuda_gen.cu new file mode 100644 index 0000000000000000000000000000000000000000..231220b203c81cd5a008b37d9e85d3eb4322e63d --- /dev/null +++ b/llm-awq/awq/kernels/csrc/quantization/gemm_cuda_gen.cu @@ -0,0 +1,298 @@ +// Inspired by NVIDIA's FasterTransformer +/* + +@article{lin2023awq, + title={AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration}, + author={Lin, Ji and Tang, Jiaming and Tang, Haotian and Yang, Shang and Dang, Xingyu and Han, Song}, + journal={arXiv}, + year={2023} +} + +*/ + +#include +#include "gemm_cuda.h" +#include +#include + + +// Pack two half values. +static inline __device__ __host__ unsigned +__pack_half2(const half x, const half y) { + unsigned v0 = *((unsigned short *)&x); + unsigned v1 = *((unsigned short *)&y); + return (v1 << 16) | v0; +} + +__device__ __forceinline__ int make_divisible(int c, int divisor){ + return (c + divisor - 1) / divisor; +} + +template +__global__ void __launch_bounds__(128) gemm_forward_4bit_cuda_m128n64k32(int split_k_iters, half* __restrict__ A, int* __restrict__ B, half* __restrict__ scaling_factors, int* zeros, int M, int IC, int OC, half* __restrict__ C) +{ + static constexpr uint32_t ZERO = 0x0; + float C_warp[64]; + __shared__ half A_shared[128 * (32 + 8)]; + __shared__ half B_shared[64 * (32 + 8)]; + + // __shared__ half scaling_factors_shared[64]; + // __shared__ half zeros_shared[64]; + + int j_factors1 = ((OC + 64 - 1) / 64); + + int blockIdx_x = 0; + int blockIdx_y = blockIdx.x % ((M + 128 - 1) / 128 * j_factors1); + int blockIdx_z = blockIdx.x / ((M + 128 - 1) / 128 * j_factors1); + + half A_shared_warp[32]; + half B_shared_warp[16]; + for (int i_0_3_init = 0; i_0_3_init < 4; ++i_0_3_init) { + for (int j_0_4_init = 0; j_0_4_init < 2; ++j_0_4_init) { + for (int i = 0; i < 8; ++i) { + C_warp[((i_0_3_init * 16) + (j_0_4_init * 8)) + i] = 0.0; + } + } + } + + static constexpr int row_stride_warp = 32 * 8 / 32; + static constexpr int row_stride_A = 4 * 32 * 8 / 32; + static constexpr int row_stride = 4 * 32 * 8 / 32; + const int make_divisible_multipler = 128 / G; + const int zeros_w = make_divisible(make_divisible(IC / G, 8), make_divisible_multipler) * make_divisible_multipler; + const int sf_w = zeros_w * 8; + + bool ld_zero_flag = (threadIdx.y * 32 + threadIdx.x) * 8 < 64; + int ld_A_row = (blockIdx_y / j_factors1 * 128 + threadIdx.y * row_stride_warp + threadIdx.x * 8 / 32); // threadIdx.y is warp_id + // bool wb_C_flag = (threadIdx.x / 4) < M; + + half* A_ptr = A + + (((int)blockIdx_y) / j_factors1 * 128 + (((int)threadIdx.y) * row_stride_warp) + ((int)threadIdx.x) / (32 / 8)) * IC + + (((int)threadIdx.x) % (32 / 8)) * 8; + + int* B_ptr = B + + ((int)threadIdx.y) * (IC / 8) * 8 + + (((int)threadIdx.x) / (32 / 8)) * (IC / 8) + + (((int)blockIdx_y) % j_factors1) * 64 * (IC / 8) + + (((int)threadIdx.x) % (32 / 8)) * 1; + +// Why * 1 in the above line? + + half* A_shared_ptr = A_shared + + ((int)threadIdx.y) * row_stride_warp * (32 + 8) + + (((int)threadIdx.x) / (32 / 8)) * (32 + 8) + + (((int)threadIdx.x) % (32 / 8) ) * 8; + + half* B_shared_ptr = B_shared + + ((int)threadIdx.y) * (row_stride / 4) * (32 + 8) + + (((int)threadIdx.x) / (32 / 8)) * (32 + 8) + + (((int)threadIdx.x) % (32 / 8)) * 8; + + + int* zeros_ptr = zeros + + ((int)threadIdx.y) * zeros_w * 8 + + (((int)threadIdx.x) / (32 / 8)) * zeros_w + + (((int)blockIdx_y) % j_factors1) * 64 * zeros_w + // this term is zero + + (((int)threadIdx.x) % (32 / 8)) / G ; + + half* scaling_factors_ptr = scaling_factors + + ((int)threadIdx.y) * sf_w * 8 + + (((int)threadIdx.x) / (32 / 8)) * sf_w + + (((int)blockIdx_y) % j_factors1) * (64) * sf_w + // this term is zero + + (((int)threadIdx.x) % (32 / 8)) * 8 / G; + + + // Haotian: TBD, check, May 29 11:46 AM PST + half* C_ptr = C + + blockIdx_z * M * OC // blockIdx_z -> split_k dim + + (((int)blockIdx_y) % j_factors1) * 64 + + (((int)threadIdx.y) / 2) * 32 + + (((int)threadIdx.x) % 4) * 2; + + // preload s.f. and zeros + int k_bound = make_divisible(IC / 32, split_k_iters); // (IC / 32 + split_k_iters - 1) / split_k_iters; + if ((k_bound - 1) * 32 + blockIdx_z >= IC) k_bound -= 1; + + // TODO (Haotian): load scales and zero points to smem + + for (int _k_0_0 = 0; _k_0_0 < k_bound; ++_k_0_0) { + int k_0_0 = _k_0_0 * split_k_iters + blockIdx_z; + __syncthreads(); + // TODO: Haotian: Here we assume M % cta_M = 0. + for (int ax0_ax1_fused_0 = 0; ax0_ax1_fused_0 < 4; ++ax0_ax1_fused_0) + { + if (ld_A_row + ax0_ax1_fused_0 * row_stride_A < M) + { + *(uint4*)(A_shared_ptr + ax0_ax1_fused_0 * row_stride_A * 40) = *(uint4*)(A_ptr + (ax0_ax1_fused_0 * row_stride_A * IC) + (k_0_0 * 32)); + } + else + { + *(uint4*)(A_shared_ptr + ax0_ax1_fused_0 * row_stride_A * 40) = make_uint4(0, 0, 0, 0); + } + } + + + int* zeros_ptr_local = zeros_ptr + k_0_0 * 32 / G / 8; + half* scaling_factors_ptr_local = scaling_factors_ptr + k_0_0 * 32 / G; + + // uint4 B_loaded_scale = make_uint4(0, 0, 0, 0); + int* B_ptr_local = B_ptr + k_0_0 * (32 / 8); + + for (int ax0_ax1_fused_0 = 0; ax0_ax1_fused_0 < 2; ++ax0_ax1_fused_0) { + + // B: 32 x 136 (128+8) float16 + // each warp: 32 x 4 + // each thr: read 32 bit -> convert to 8xFP16 (a UINT4) -> scale and minus zero -> WB UINT4 + // row stride in shared memory: (NWARPS * 32 * 8 / cta_N) + int B_loaded_current = *(B_ptr_local + ax0_ax1_fused_0 * row_stride * (IC / 8)); + int zeros_loaded = *(zeros_ptr_local + ax0_ax1_fused_0 * row_stride * zeros_w); + zeros_loaded >>= ((k_0_0 * 32 / G) % 8) * 4; + float current_zeros = (float)(zeros_loaded & 0xF); + half scaling_factors_loaded = *(scaling_factors_ptr_local + ax0_ax1_fused_0 * row_stride * sf_w); + half B_loaded_fp16[8]; + #pragma unroll + for (int ic_1 = 0; ic_1 < 8; ic_1++){ + float current_single_weight_fp = (float)(B_loaded_current & 0xF); + half dequantized_weight = __float2half(__half2float(scaling_factors_loaded) * (current_single_weight_fp - current_zeros)); + B_loaded_current = B_loaded_current >> 4; + B_loaded_fp16[ic_1] = dequantized_weight; + } + // write back + *(uint4*)(B_shared_ptr + ax0_ax1_fused_0 * row_stride * (32 + 8)) = *reinterpret_cast(B_loaded_fp16); + } + __syncthreads(); + for (int k_0_1 = 0; k_0_1 < 2; ++k_0_1) { + for (int ax0_0 = 0; ax0_0 < 4; ++ax0_0) { + { + unsigned int addr; + __asm__ __volatile__( + "{ .reg .u64 addr; cvta.to.shared.u64 addr, %1; cvt.u32.u64 %0, addr; }\n" + : "=r"(addr) + : "l"((void *)((&(A_shared[((((((int)threadIdx.y) & 1) * 2560) + (ax0_0 * 640)) + (k_0_1 * 16))])) + (((((int)threadIdx.x) & 15) * 40) + ((((int)threadIdx.x) >> 4) * 8)))) + ); + __asm__ __volatile__( + "ldmatrix.sync.aligned.m8n8.x4.shared.b16" + "{%0, %1, %2, %3}, [%4];\n" + : "=r"(((unsigned *)(A_shared_warp + (ax0_0 * 8)))[0]), "=r"(((unsigned *)(A_shared_warp + (ax0_0 * 8)))[1]), "=r"(((unsigned *)(A_shared_warp + (ax0_0 * 8)))[2]), "=r"(((unsigned *)(A_shared_warp + (ax0_0 * 8)))[3]) + : "r"(addr) + ); + } + } + + for (int ax0_0_1 = 0; ax0_0_1 < 2; ++ax0_0_1) { + { + unsigned int addr; + __asm__ __volatile__( + "{ .reg .u64 addr; cvta.to.shared.u64 addr, %1; cvt.u32.u64 %0, addr; }\n" + : "=r"(addr) + : "l"((void *)((&(B_shared[((((((int)threadIdx.y) >> 1) * 1280) + (ax0_0_1 * 640)) + (k_0_1 * 16))])) + ((((((int)threadIdx.x) >> 4) * 320) + ((((int)threadIdx.x) & 7) * 40)) + (((((int)threadIdx.x) & 15) >> 3) * 8)))) + ); + __asm__ __volatile__( + "ldmatrix.sync.aligned.m8n8.x4.shared.b16" + "{%0, %1, %2, %3}, [%4];\n" + : "=r"(((unsigned *)(B_shared_warp + (ax0_0_1 * 8)))[0]), "=r"(((unsigned *)(B_shared_warp + (ax0_0_1 * 8)))[1]), "=r"(((unsigned *)(B_shared_warp + (ax0_0_1 * 8)))[2]), "=r"(((unsigned *)(B_shared_warp + (ax0_0_1 * 8)))[3]) + : "r"(addr) + ); + } + } + + for (int i_0_3 = 0; i_0_3 < 4; ++i_0_3) { + for (int j_0_4 = 0; j_0_4 < 2; ++j_0_4) { + + { + __asm__ __volatile__( + "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32" + "{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%10, %11, %12, %13};\n" + : "=f"(((float *)(C_warp + ((i_0_3 * 16) + (j_0_4 * 8))))[0]), "=f"(((float *)(C_warp + ((i_0_3 * 16) + (j_0_4 * 8))))[1]), "=f"(((float *)(C_warp + ((i_0_3 * 16) + (j_0_4 * 8))))[2]), "=f"(((float *)(C_warp + ((i_0_3 * 16) + (j_0_4 * 8))))[3]) + : "r"(((unsigned *)(A_shared_warp + (i_0_3 * 8)))[0]), "r"(((unsigned *)(A_shared_warp + (i_0_3 * 8)))[1]), "r"(((unsigned *)(A_shared_warp + (i_0_3 * 8)))[2]), "r"(((unsigned *)(A_shared_warp + (i_0_3 * 8)))[3]), "r"(((unsigned *)(B_shared_warp + (j_0_4 * 8)))[0]), "r"(((unsigned *)(B_shared_warp + (j_0_4 * 8)))[1]), "f"(((float *)(C_warp + ((i_0_3 * 16) + (j_0_4 * 8))))[0]), "f"(((float *)(C_warp + ((i_0_3 * 16) + (j_0_4 * 8))))[1]), "f"(((float *)(C_warp + ((i_0_3 * 16) + (j_0_4 * 8))))[2]), "f"(((float *)(C_warp + ((i_0_3 * 16) + (j_0_4 * 8))))[3])); + } + + { + __asm__ __volatile__( + "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32" + "{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%10, %11, %12, %13};\n" + : "=f"(((float *)(C_warp + (((i_0_3 * 16) + (j_0_4 * 8)) + 4)))[0]), "=f"(((float *)(C_warp + (((i_0_3 * 16) + (j_0_4 * 8)) + 4)))[1]), "=f"(((float *)(C_warp + (((i_0_3 * 16) + (j_0_4 * 8)) + 4)))[2]), "=f"(((float *)(C_warp + (((i_0_3 * 16) + (j_0_4 * 8)) + 4)))[3]) + : "r"(((unsigned *)(A_shared_warp + (i_0_3 * 8)))[0]), "r"(((unsigned *)(A_shared_warp + (i_0_3 * 8)))[1]), "r"(((unsigned *)(A_shared_warp + (i_0_3 * 8)))[2]), "r"(((unsigned *)(A_shared_warp + (i_0_3 * 8)))[3]), "r"(((unsigned *)(B_shared_warp + ((j_0_4 * 8) + 4)))[0]), "r"(((unsigned *)(B_shared_warp + ((j_0_4 * 8) + 4)))[1]), "f"(((float *)(C_warp + (((i_0_3 * 16) + (j_0_4 * 8)) + 4)))[0]), "f"(((float *)(C_warp + (((i_0_3 * 16) + (j_0_4 * 8)) + 4)))[1]), "f"(((float *)(C_warp + (((i_0_3 * 16) + (j_0_4 * 8)) + 4)))[2]), "f"(((float *)(C_warp + (((i_0_3 * 16) + (j_0_4 * 8)) + 4)))[3])); + } + } + } + } + } + +// Haotian: Here (May 29 11:46AM PST) +// TODO: Shang: Hoist loop invariance. + for (int ax0_0_2 = 0; ax0_0_2 < 4; ++ax0_0_2) { + for (int ax1_0 = 0; ax1_0 < 2; ++ax1_0) { + for (int local_id = 0; local_id < 8; ++local_id) { + int row_offset = (((int)blockIdx_y) / j_factors1) * 128 + (threadIdx.y % 2) * 64 + ax0_0_2 * 16 + (local_id % 4) / 2 * 8 + ((int)threadIdx.x) / 4; + if (row_offset < M) + { + *(C_ptr + ax1_0 * 16 + row_offset * OC + (local_id / 4) * 8 + local_id % 2) = __float2half(C_warp[(ax0_0_2 * 16) + (ax1_0 * 8) + local_id]); + } + } + } + } +} + +// in_feats: M, IC [float16] +// kernel: IC, OC // 8 [int32] -> cast to IC, OC [uint4b] +// scaling_factors: IC // G, OC [float16] +// zeros: IC // G, OC // 8 [int32] -> cast to IC // G, OC [uint4b] +// assume that batch_size < 16 for now + +torch::Tensor gemm_forward_cuda( + torch::Tensor _in_feats, + torch::Tensor _kernel, + torch::Tensor _scaling_factors, + torch::Tensor _zeros, + int group_size, + int split_k_iters) +{ + int num_in_feats = _in_feats.size(0); + int num_in_channels = _in_feats.size(1); + const at::cuda::OptionalCUDAGuard device_guard(device_of(_in_feats)); + + auto options = torch::TensorOptions().dtype(_in_feats.dtype()).device(_in_feats.device()); + // for int4, need _kernel.size(1) * 8 + at::Tensor _out_feats = torch::empty({split_k_iters, num_in_feats, _kernel.size(0)}, options); + int num_out_feats = _out_feats.size(-2); + int num_out_channels = _out_feats.size(-1); + + auto in_feats = reinterpret_cast(_in_feats.data_ptr()); + auto kernel = reinterpret_cast(_kernel.data_ptr()); + auto out_feats = reinterpret_cast(_out_feats.data_ptr()); + auto scaling_factors = reinterpret_cast(_scaling_factors.data_ptr()); + auto zeros = reinterpret_cast(_zeros.data_ptr()); + + // blockIdx_x: i_factors[0] * j_factors[0] + // blockIdx_y: i_factors[1] * j_factors[1] + + if (num_out_channels % 64 != 0) + throw std::invalid_argument("OC is not multiple of cta_N = 64"); + if (num_out_channels % 8 != 0) + throw std::invalid_argument("OC is not multiple of pack_num = 8"); + int j_factors1 = num_out_channels / 64 / 1; + dim3 num_blocks((num_out_feats + 128 - 1) / 128 * j_factors1 * split_k_iters); + + // threadIdx.x: 32 + // threadIdx.y: i_factors[2] * j_factors[2] + dim3 threads_per_block(32, 4); + if (group_size == 128) + { + gemm_forward_4bit_cuda_m128n64k32<128><<>>( + split_k_iters, in_feats, kernel, scaling_factors, zeros, num_in_feats, num_in_channels, num_out_channels, out_feats); + } + else if (group_size == 64) + { + gemm_forward_4bit_cuda_m128n64k32<64><<>>( + split_k_iters, in_feats, kernel, scaling_factors, zeros, num_in_feats, num_in_channels, num_out_channels, out_feats); + } + else + { + throw std::invalid_argument("Group size temporarily not supported."); + } + return _out_feats.sum(0); +} + diff --git a/llm-awq/awq/kernels/csrc/quantization/gemv_cuda.cu b/llm-awq/awq/kernels/csrc/quantization/gemv_cuda.cu new file mode 100644 index 0000000000000000000000000000000000000000..3a55e662e82d77d5c9322fef32df0f24bc2bb911 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/quantization/gemv_cuda.cu @@ -0,0 +1,247 @@ +// Inspired by https://github.com/ankan-ban/llama_cu_awq +/* + +@article{lin2023awq, + title={AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration}, + author={Lin, Ji and Tang, Jiaming and Tang, Haotian and Yang, Shang and Dang, Xingyu and Han, Song}, + journal={arXiv}, + year={2023} +} + +*/ + +#include +#include +#include +#include "gemv_cuda.h" +#define VECTORIZE_FACTOR 8 +#define Q_VECTORIZE_FACTOR 8 +#define PACK_FACTOR 8 +#define WARP_SIZE 32 + + +// Reduce sum within the warp using the tree reduction algorithm. +__device__ __forceinline__ float warp_reduce_sum(float sum) { + #pragma unroll + for(int i = 4; i >= 0; i--){ + sum += __shfl_down_sync(0xffffffff, sum, 1<(zeros + oc_idx * zeros_w + packed_group_idx * 2); + uint32_t packed_weights[4]; + // use float4 to load weights, each thread load 32 int4 numbers (1 x float4) + *((float4*)(packed_weights)) = *((float4*)(weight + oc_idx * weight_w + packed_group_idx * (WARP_SIZE * 4) + threadIdx.x * 4)); + // load scaling factors + // g64: two threads -> 64 numbers -> 1 group; 1 warp = 16 groups. + float scaling_factor = __half2float(scaling_factors[oc_idx * sf_w + packed_group_idx * 16 + (threadIdx.x / 2)]); + float current_zeros = (float)((packed_zeros >> (threadIdx.x / 2 * 4)) & 0xF); + int inputs_ptr_delta = packed_group_idx * WARP_SIZE * 4 + threadIdx.x * 4; + const float4* inputs_ptr = inputs + inputs_ptr_delta; + // multiply 32 weights with 32 inputs + #pragma unroll + for (int ic_0 = 0; ic_0 < 4; ic_0++){ + // iterate over different uint32_t packed_weights in this loop + uint32_t current_packed_weight = packed_weights[ic_0]; + half packed_inputs[PACK_FACTOR]; + // each thread load 8 inputs, starting index is packed_group_idx * 128 * 8 (because each iter loads 128*8) + if (inputs_ptr_delta + ic_0 < IC / PACK_FACTOR) { + *((float4*)packed_inputs) = *(inputs_ptr + ic_0); + #pragma unroll + for (int ic_1 = 0; ic_1 < PACK_FACTOR; ic_1++){ + // iterate over 8 numbers packed within each uint32_t number + float current_single_weight_fp = (float)(current_packed_weight & 0xF); + float dequantized_weight = scaling_factor * (current_single_weight_fp - current_zeros); + //if(blockIdx.x == 0 && blockIdx.y == 0 && threadIdx.x == 0 && threadIdx.y == 0 && ic_0 == 0 && ic_1 == 0 && packed_group_idx == 0) printf("%f %f %f %f %X %X\n", dequantized_weight, current_single_weight_fp, scaling_factor, current_zeros, current_packed_weight, packed_zeros); + psum += dequantized_weight * __half2float(packed_inputs[ic_1]); + current_packed_weight = current_packed_weight >> 4; + } + } + } + } + psum = warp_reduce_sum(psum); + if (threadIdx.x == 0) { + outputs[oc_idx] = __float2half(psum); + } +} + + +/* +Computes GEMV (group_size = 128). + +Args: + inputs: vector of shape [batch_size, IC]; + weight: matrix of shape [OC, IC / 8]; + output: vector of shape [OC]; + zeros: matrix of shape [OC, IC / group_size / 8]; + scaling_factors: matrix of shape [OC, IC / group_size]; + +Notes: + One cannot infer group_size from the shape of scaling factors. + the second dimension is rounded up to a multiple of PACK_FACTOR. +*/ +__global__ void gemv_kernel_g128( + const float4* _inputs, const uint32_t* weight, const uint32_t* zeros, const half* scaling_factors, half* _outputs, + const int IC, const int OC){ + const int group_size = 128; + float psum = 0; + const int batch_idx = blockIdx.z; + const int oc_idx = blockIdx.y * blockDim.y + threadIdx.y; + const float4* inputs = _inputs + batch_idx * IC / PACK_FACTOR; + half* outputs = _outputs + batch_idx * OC; + const int num_groups_packed = make_divisible(IC / group_size, PACK_FACTOR); + const int weight_w = IC / PACK_FACTOR; + // TODO (Haotian): zeros_w is incorrect, after fixing we got misaligned address + const int zeros_w = make_divisible(IC / group_size, PACK_FACTOR); + // consistent with input shape + const int sf_w = make_divisible(IC / group_size, PACK_FACTOR) * PACK_FACTOR; + //if(blockIdx.x == 0 && blockIdx.y == 0 && threadIdx.x == 0 && threadIdx.y == 0) printf("%d %d %d %d\n", IC, group_size, PACK_FACTOR, zeros_w); + // tile size: 4 OC x 1024 IC per iter + for(int packed_group_idx = 0; packed_group_idx < num_groups_packed; packed_group_idx++){ + // 1024 numbers in one iteration across warp. Need 1024 / group_size zeros. + uint32_t packed_zeros = *(zeros + oc_idx * zeros_w + packed_group_idx); + uint32_t packed_weights[4]; + // use float4 to load weights, each thread load 32 int4 numbers (1 x float4) + *((float4*)(packed_weights)) = *((float4*)(weight + oc_idx * weight_w + packed_group_idx * (WARP_SIZE * 4) + threadIdx.x * 4)); + // load scaling factors + // g128: four threads -> 128 numbers -> 1 group; 1 warp = 8 groups. + float scaling_factor = __half2float(scaling_factors[oc_idx * sf_w + packed_group_idx * 8 + (threadIdx.x / 4)]); + float current_zeros = (float)((packed_zeros >> (threadIdx.x / 4 * 4)) & 0xF); + int inputs_ptr_delta = packed_group_idx * WARP_SIZE * 4 + threadIdx.x * 4; + const float4* inputs_ptr = inputs + inputs_ptr_delta; + // multiply 32 weights with 32 inputs + #pragma unroll + for (int ic_0 = 0; ic_0 < 4; ic_0++){ + // iterate over different uint32_t packed_weights in this loop + uint32_t current_packed_weight = packed_weights[ic_0]; + half packed_inputs[PACK_FACTOR]; + // each thread load 8 inputs, starting index is packed_group_idx * 128 * 8 (because each iter loads 128*8) + if (inputs_ptr_delta + ic_0 < IC / PACK_FACTOR) { + *((float4*)packed_inputs) = *(inputs_ptr + ic_0); + #pragma unroll + for (int ic_1 = 0; ic_1 < PACK_FACTOR; ic_1++){ + // iterate over 8 numbers packed within each uint32_t number + float current_single_weight_fp = (float)(current_packed_weight & 0xF); + float dequantized_weight = scaling_factor * (current_single_weight_fp - current_zeros); + //if(blockIdx.x == 0 && blockIdx.y == 0 && threadIdx.x == 0 && threadIdx.y == 0 && ic_0 == 0 && ic_1 == 0 && packed_group_idx == 0) printf("%f %f %f %f %X %X\n", dequantized_weight, current_single_weight_fp, scaling_factor, current_zeros, current_packed_weight, packed_zeros); + psum += dequantized_weight * __half2float(packed_inputs[ic_1]); + current_packed_weight = current_packed_weight >> 4; + } + } + } + } + psum = warp_reduce_sum(psum); + if (threadIdx.x == 0) { + outputs[oc_idx] = __float2half(psum); + } +} + + +/* +Computes GEMV (PyTorch interface). + +Args: + _in_feats: tensor of shape [B, IC]; + _kernel: int tensor of shape [OC, IC // 8]; + _zeros: int tensor of shape [OC, IC // G // 8]; + _scaling_factors: tensor of shape [OC, IC // G]; + blockDim_x: size of thread block, dimension x, where blockDim_x * workload_per_thread = IC; + blockDim_y: size of thread block, dimension y, where blockDim_y * gridDim_y = OC; + +Returns: + out_feats: tensor of shape [B, OC]; +*/ +torch::Tensor gemv_forward_cuda( + torch::Tensor _in_feats, + torch::Tensor _kernel, + torch::Tensor _scaling_factors, + torch::Tensor _zeros, + int group_size) +{ + int num_in_feats = _in_feats.size(0); + int num_in_channels = _in_feats.size(1); + // int kernel_volume = _out_in_map.size(1); + auto in_feats = reinterpret_cast(_in_feats.data_ptr()); + auto kernel = reinterpret_cast(_kernel.data_ptr()); + auto zeros = reinterpret_cast(_zeros.data_ptr()); + auto scaling_factors = reinterpret_cast(_scaling_factors.data_ptr()); + // auto out_in_map = _out_in_map.data_ptr(); + auto options = + torch::TensorOptions().dtype(_in_feats.dtype()).device(_in_feats.device()); + // kernel is [OC, IC] + at::Tensor _out_feats = torch::empty({num_in_feats, _kernel.size(0)}, options); + int num_out_feats = _out_feats.size(-2); + int num_out_channels = _out_feats.size(-1); + auto out_feats = reinterpret_cast(_out_feats.data_ptr()); + int blockDim_z = num_out_feats; + dim3 num_blocks(1, num_out_channels / 4, num_out_feats); + dim3 num_threads(32, 4); + if (group_size == 64) + { + gemv_kernel_g64<<>>( + // pointers + in_feats, kernel, zeros, scaling_factors, out_feats, + // constants + num_in_channels, num_out_channels + ); + } + else if (group_size == 128) + { + gemv_kernel_g128<<>>( + // pointers + in_feats, kernel, zeros, scaling_factors, out_feats, + // constants + num_in_channels, num_out_channels + ); + } + return _out_feats; +;} + diff --git a/llm-awq/awq/kernels/csrc/quantization/gemv_cuda.h b/llm-awq/awq/kernels/csrc/quantization/gemv_cuda.h new file mode 100644 index 0000000000000000000000000000000000000000..748abc5d1bcef5cca0c056c6e1e2279c3bb262a9 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/quantization/gemv_cuda.h @@ -0,0 +1,9 @@ +#pragma once +#include + +torch::Tensor gemv_forward_cuda( + torch::Tensor _in_feats, + torch::Tensor _kernel, + torch::Tensor _scaling_factors, + torch::Tensor _zeros, + int group_size); diff --git a/llm-awq/awq/kernels/csrc/quantization_new/dequantize.cuh b/llm-awq/awq/kernels/csrc/quantization_new/dequantize.cuh new file mode 100644 index 0000000000000000000000000000000000000000..9917ec94911ca106844858ea855f4fc15902a2a7 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/quantization_new/dequantize.cuh @@ -0,0 +1,123 @@ +/* +Modified from NVIDIA FasterTransformer: https://github.com/NVIDIA/FasterTransformer/blob/main/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h + +@article{lin2023awq, + title={AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration}, + author={Lin, Ji and Tang, Jiaming and Tang, Haotian and Yang, Shang and Dang, Xingyu and Han, Song}, + journal={arXiv}, + year={2023} +} +*/ +#include +#include +#pragma once + +template +__inline__ __device__ void dequantize_s4_to_fp16x2(half2 const &source, uint4 *result); + +template <> +__inline__ __device__ void dequantize_s4_to_fp16x2(half2 const &source, uint4 *result) +{ + // uint4 result; + + uint32_t *h = reinterpret_cast(result); + uint32_t const i4s = reinterpret_cast(source); + + // First, we extract the i4s and construct an intermediate fp16 number. + constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa; + constexpr uint32_t BOTTOM_MASK = 0x000f000f; + constexpr uint32_t TOP_MASK = 0x00f000f0; + constexpr uint32_t I4s_TO_F16s_MAGIC_NUM = 0x64006400; + + // Note that the entire sequence only requires 1 shift instruction. This is thanks to the register packing + // format and the fact that we force our integers to be unsigned, and account for this in the fp16 subtractions. + // In addition, I exploit the fact that sub and fma have the same throughput in order to convert elt_23 and + // elt_67 to fp16 without having to shift them to the bottom bits before hand. + + // Shift right by 8 to now consider elt_45 and elt_67. Issue first to hide RAW dependency if we issue + // immediately before required. + const uint32_t top_i4s = i4s >> 8; + // Extract elt_01 - (i4s & 0x000f000f) | 0x64006400 + asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(h[0]) + : "r"(i4s), "n"(BOTTOM_MASK), "n"(I4s_TO_F16s_MAGIC_NUM), "n"(immLut)); + // Extract elt_23 (i4s & 0x00f000f0) | 0x64006400 + asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(h[1]) + : "r"(i4s), "n"(TOP_MASK), "n"(I4s_TO_F16s_MAGIC_NUM), "n"(immLut)); + // Extract elt_45 (top_i4s & 0x000f000f) | 0x64006400 + asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(h[2]) + : "r"(top_i4s), "n"(BOTTOM_MASK), "n"(I4s_TO_F16s_MAGIC_NUM), "n"(immLut)); + // Extract elt_67 (top_i4s & 0x00f000f0) | 0x64006400 + asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(h[3]) + : "r"(top_i4s), "n"(TOP_MASK), "n"(I4s_TO_F16s_MAGIC_NUM), "n"(immLut)); + + // I use inline PTX below because I am not sure if the compiler will emit float2half instructions if I use the + // half2 ctor. In this case, I chose performance reliability over code readability. + + // This is the half2 {1032, 1032} represented as an integer. + // static constexpr uint32_t FP16_TOP_MAGIC_NUM = 0x64086408; + // Haotian: subtract {1024, 1024} instead, we do not need to map to [-8, 7] + static constexpr uint32_t FP16_TOP_MAGIC_NUM = 0x64006400; + // This is the half2 {1 / 16, 1 / 16} represented as an integer. + static constexpr uint32_t ONE_SIXTEENTH = 0x2c002c00; + // This is the half2 {-72, -72} represented as an integer. + // static constexpr uint32_t NEG_72 = 0xd480d480; + // Haotian: Let's use {-64, -64}. + static constexpr uint32_t NEG_64 = 0xd400d400; + + // Finally, we construct the output numbers. + // Convert elt_01 + asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(h[0]) : "r"(h[0]), "r"(FP16_TOP_MAGIC_NUM)); + // Convert elt_23 + asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n" : "=r"(h[1]) : "r"(h[1]), "r"(ONE_SIXTEENTH), "r"(NEG_64)); + // Convert elt_45 + asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(h[2]) : "r"(h[2]), "r"(FP16_TOP_MAGIC_NUM)); + // Convert elt_67 + asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n" : "=r"(h[3]) : "r"(h[3]), "r"(ONE_SIXTEENTH), "r"(NEG_64)); +} + +template <> +__inline__ __device__ void dequantize_s4_to_fp16x2(half2 const &source, uint4 *result) +{ + // uint4 result; + + uint32_t *h = reinterpret_cast(result); + uint32_t const i4s = reinterpret_cast(source); + + // First, we extract the i4s and construct an intermediate bf16 number. + constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa; + constexpr uint32_t BOTTOM_MASK = 0x000f000f; + constexpr uint32_t I4s_TO_BF16s_MAGIC_NUM = 0x43004300; + + // Shift right by 4, 8, 12 to consider elt_23, elt_45 and elt_67. + const uint32_t i4s1 = i4s >> 4; + const uint32_t i4s2 = i4s >> 8; + const uint32_t i4s3 = i4s >> 12; + // Extract elt_01 - (i4s & 0x000f000f) | 0x43004300 + asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(h[0]) + : "r"(i4s), "n"(BOTTOM_MASK), "n"(I4s_TO_BF16s_MAGIC_NUM), "n"(immLut)); + // Extract elt_23 - (i4s & 0x000f000f) | 0x43004300 + asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(h[1]) + : "r"(i4s1), "n"(BOTTOM_MASK), "n"(I4s_TO_BF16s_MAGIC_NUM), "n"(immLut)); + // Extract elt_45 - (i4s & 0x000f000f) | 0x43004300 + asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(h[2]) + : "r"(i4s2), "n"(BOTTOM_MASK), "n"(I4s_TO_BF16s_MAGIC_NUM), "n"(immLut)); + // Extract elt_67 - (i4s & 0x000f000f) | 0x43004300 + asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(h[3]) + : "r"(i4s3), "n"(BOTTOM_MASK), "n"(I4s_TO_BF16s_MAGIC_NUM), "n"(immLut)); + + // This is the nv_bfloat162 {128, 128} represented as an integer + static constexpr uint32_t BF16_TOP_MAGIC_NUM = 0x43004300; + + reinterpret_cast<__nv_bfloat162*>(h)[0] = __hsub2(reinterpret_cast<__nv_bfloat162*>(h)[0], reinterpret_cast(BF16_TOP_MAGIC_NUM)); + reinterpret_cast<__nv_bfloat162*>(h)[1] = __hsub2(reinterpret_cast<__nv_bfloat162*>(h)[1], reinterpret_cast(BF16_TOP_MAGIC_NUM)); + reinterpret_cast<__nv_bfloat162*>(h)[2] = __hsub2(reinterpret_cast<__nv_bfloat162*>(h)[2], reinterpret_cast(BF16_TOP_MAGIC_NUM)); + reinterpret_cast<__nv_bfloat162*>(h)[3] = __hsub2(reinterpret_cast<__nv_bfloat162*>(h)[3], reinterpret_cast(BF16_TOP_MAGIC_NUM)); +} \ No newline at end of file diff --git a/llm-awq/awq/kernels/csrc/quantization_new/dispatch_utils.cuh b/llm-awq/awq/kernels/csrc/quantization_new/dispatch_utils.cuh new file mode 100644 index 0000000000000000000000000000000000000000..cb7773838cad0c767e6bb4aca1668ee2cf7e2154 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/quantization_new/dispatch_utils.cuh @@ -0,0 +1,18 @@ +#pragma once +#include +#include +#include +#include + +#define DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP16(pytorch_dtype, c_type, ...) \ + if (pytorch_dtype == at::ScalarType::Half) { \ + using c_type = half; \ + __VA_ARGS__ \ + } else if (pytorch_dtype == at::ScalarType::BFloat16) { \ + using c_type = nv_bfloat16; \ + __VA_ARGS__ \ + } else { \ + std::ostringstream oss; \ + oss << __PRETTY_FUNCTION__ << " failed to dispatch data type " << pytorch_dtype; \ + TORCH_CHECK(false, oss.str()); \ + } diff --git a/llm-awq/awq/kernels/csrc/quantization_new/gemm/gemm_cuda.cu b/llm-awq/awq/kernels/csrc/quantization_new/gemm/gemm_cuda.cu new file mode 100644 index 0000000000000000000000000000000000000000..d4d297da5a919d1cc7871a57b49924fdc36b39b9 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/quantization_new/gemm/gemm_cuda.cu @@ -0,0 +1,1236 @@ +#include +#include "semaphore.h" +#include "gemm_cuda.h" +#include "../dequantize.cuh" +#include "../dispatch_utils.cuh" +#include +#include + +#define kInterleave 4 +#define OP_M 16 +#define OP_N 8 +#define OP_K 16 +#define INTRIN_M 16 +#define INTRIN_N 16 +#define INTRIN_K 16 +#define WARP_SIZE 32 +#define SMEM_PAD_A 0 +#define SMEM_PAD_B 0 +#define PACK_SIZE 8 +#if (__CUDACC_VER_MAJOR__ >= 11) && (__CUDACC_VER_MINOR__ >= 4) +#define L2_CACHEHINT(size) ".L2::" #size "B" +#else +#define L2_CACHEHINT(size) +#endif + +#define KERNEL_LAUNCH_CODE \ + int num_mn_tiles = (num_in_feats + CTA_M - 1) / CTA_M * (num_out_channels + CTA_N - 1) / CTA_N; \ + torch::Tensor _semaphores = torch::empty({num_mn_tiles}, options_int); \ + auto semaphores = reinterpret_cast(_semaphores.data_ptr()); \ + constexpr int NUM_WARPS = (CTA_M / WARP_M) * (CTA_N / WARP_N) * (CTA_K / WARP_K); \ + constexpr int SCALES_SMEM_SIZE = (G >= CTA_K) ? (CTA_N / (G / CTA_K) * STAGES * 2) : (CTA_N * (CTA_K / G) * STAGES * 2); \ + constexpr int kSmemByteSize = (CTA_M * (CTA_K + SMEM_PAD_A) + CTA_N * (CTA_K + SMEM_PAD_B) / kInterleave + SCALES_SMEM_SIZE) * STAGES * sizeof(ctype); \ + if (kSmemByteSize >= 99 * 1024) \ + { \ + printf("This kernel requires %d Bytes of shared memory, which exceeds device limit.\n", kSmemByteSize); \ + return _out_feats; \ + } \ + int j_factors1 = num_out_channels / CTA_N / 1; \ + dim3 num_blocks((num_out_feats + CTA_M - 1) / CTA_M * j_factors1 * SPLITK); \ + dim3 threads_per_block(WARP_SIZE, NUM_WARPS); \ + auto kernel_func = gemm_w4a16_T1; \ + cudaFuncSetAttribute(kernel_func, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemByteSize); \ + kernel_func<<>>( \ + in_feats, kernel, scales, zeros, out_feats, semaphores, num_in_feats, num_out_channels, num_in_channels); + +template +__inline__ __host__ __device__ int get_log_tile(int n) +{ + if (N >= 8 && n >= 6) + return 3; + else if (N >= 4 && n >= 3) + return 2; + else if (N >= 2 && n >= 2) + return 1; + else + return 0; +} + +__inline__ __device__ uint2 get_block_idx_mapping(int blockIdx_x, int blockIdx_y, int log_tile) +{ + return make_uint2((blockIdx_x >> log_tile), (blockIdx_y << log_tile) + ((blockIdx_x) & ((1 << (log_tile)) - 1))); +} + +template +__device__ void sync_slice(int slice_id) +{ + if constexpr (SLICES == 1) + { + __syncthreads(); + } + else + { + constexpr int SLICE_GROUP = (SLICES + 7) / 8; + constexpr uint32_t num_threads = NUM_WARPS_MN * WARP_SIZE; + const uint32_t barrier_id = slice_id / SLICE_GROUP + 1; + asm volatile("bar.sync %0, %1;" : : "r"(barrier_id), "n"(num_threads)); + } +} + +__inline__ __device__ uint32_t cast_smem_ptr_to_uint(void const *const ptr) +{ + uint32_t smem_int_ptr; + + asm("{.reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }\n" + : "=r"(smem_int_ptr) + : "l"(ptr)); + + return smem_int_ptr; +} + +template +__inline__ __device__ void ldmatrix_m8n8_x4_b16(T *shared_warp, int ax0_0, uint32_t addr) +{ + __asm__ __volatile__( + "ldmatrix.sync.aligned.m8n8.x4.shared.b16" + "{%0, %1, %2, %3}, [%4];" + : "=r"(((unsigned *)(shared_warp + (ax0_0 * 8)))[0]), "=r"(((unsigned *)(shared_warp + (ax0_0 * 8)))[1]), "=r"(((unsigned *)(shared_warp + (ax0_0 * 8)))[2]), "=r"(((unsigned *)(shared_warp + (ax0_0 * 8)))[3]) + : "r"(addr)); +} + +template +__inline__ __device__ void ldmatrix_m8n8_x4_trans_b16(T *shared_warp, int ax0_0, uint32_t addr) +{ + __asm__ __volatile__( + "ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16" + "{%0, %1, %2, %3}, [%4];" + : "=r"(((unsigned *)(shared_warp + (ax0_0 * 8)))[0]), "=r"(((unsigned *)(shared_warp + (ax0_0 * 8)))[1]), "=r"(((unsigned *)(shared_warp + (ax0_0 * 8)))[2]), "=r"(((unsigned *)(shared_warp + (ax0_0 * 8)))[3]) + : "r"(addr)); +} + +__inline__ __device__ void cp_async_cg_A(uint32_t smem_int_ptr, const uint4 *__restrict__ src, bool mask) +{ + const int cp_size = 16; + asm volatile("{" + " .reg .pred p;" + " setp.ne.b32 p, %0, 0;" + " @p cp.async.cg.shared.global" L2_CACHEHINT(128) " [%1], [%2], %3;" + "}" ::"r"((int)mask), + "r"(smem_int_ptr), + "l"(src), + "n"(cp_size)); +} + +__device__ __inline__ void mma_m16n8k16_f16f16f16(half *C_warp, half *A_shared_warp, half *B_shared_warp) +{ + __asm__ __volatile__( + "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16" + "{%0, %1}, {%2, %3, %4, %5}, {%6, %7}, {%8, %9};" + : "=r"(((unsigned *)C_warp)[0]), "=r"(((unsigned *)C_warp)[1]) + : "r"(((unsigned *)A_shared_warp)[0]), "r"(((unsigned *)A_shared_warp)[1]), "r"(((unsigned *)A_shared_warp)[2]), "r"(((unsigned *)A_shared_warp)[3]), "r"(((unsigned *)B_shared_warp)[0]), "r"(((unsigned *)B_shared_warp)[1]), "r"(((unsigned *)C_warp)[0]), "r"(((unsigned *)C_warp)[1])); +} + +__device__ __inline__ void mma_m16n8k16_bf16bf16f32(float *C_warp, nv_bfloat16 *A_shared_warp, nv_bfloat16 *B_shared_warp) +{ + + __asm__ __volatile__( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32" + "{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%10, %11, %12, %13};" + : "=f"(C_warp[0]), "=f"(C_warp[1]), "=f"(C_warp[2]), "=f"(C_warp[3]) + : "r"(((unsigned *)A_shared_warp)[0]), "r"(((unsigned *)A_shared_warp)[1]), "r"(((unsigned *)A_shared_warp)[2]), "r"(((unsigned *)A_shared_warp)[3]), "r"(((unsigned *)B_shared_warp)[0]), "r"(((unsigned *)B_shared_warp)[1]), "f"(C_warp[0]), "f"(C_warp[1]), "f"(C_warp[2]), "f"(C_warp[3])); +} + +template +__device__ __inline__ void global_to_share_one_stage_A(T *src, T *dst, int global_nrows, int global_ncols, int cta_offset_m, int cta_offset_n, int cta_offset_k, int global_iter_k, int shared_iter_k, bool mask) +{ + constexpr int threads_needed = (CTA_M * CTA_K) / PACK_SIZE / SHARED_K_ITERS; + constexpr int threads_used = threads_needed < CTA_SIZE ? threads_needed : CTA_SIZE; + constexpr int total_global_iters = (CTA_M * CTA_K) / PACK_SIZE / threads_used; + constexpr int partial_global_iters = (total_global_iters + SHARED_K_ITERS - 1) / SHARED_K_ITERS; + constexpr int cta_step_m_or_n = (threads_used * PACK_SIZE) / CTA_K; + constexpr int warp_step_m_or_n = (WARP_SIZE * PACK_SIZE) / CTA_K; + constexpr int threads_per_row = CTA_K / PACK_SIZE; + constexpr int kSmemCol = CTA_K + SMEM_PAD_A; + bool local_mask = mask & (threadIdx.y * WARP_SIZE + threadIdx.x < threads_used); + int ld_col = (threadIdx.x % threads_per_row); +#pragma unroll + for (int _global_iter = 0; _global_iter < partial_global_iters; ++_global_iter) + { + int global_iter = shared_iter_k * partial_global_iters + _global_iter; + int ld_row = global_iter * cta_step_m_or_n + threadIdx.y * warp_step_m_or_n + (threadIdx.x / threads_per_row); + int ld_col_swizzled = (ld_col ^ (ld_row) & 7) * PACK_SIZE; + void *dst_ptr = (void *)(dst + ld_row * kSmemCol + ld_col_swizzled); + uint4 *src_ptr = (uint4 *)(src + (ld_row + cta_offset_m) * global_ncols + ld_col * PACK_SIZE + global_iter_k * CTA_K + cta_offset_k); // cta_offset_m * global_ncols + global_iter * cta_step_m_or_n * global_ncols + threadIdx.y * warp_step_m_or_n * global_ncols + (threadIdx.x / threads_per_row) * global_ncols + global_iter_k * CTA_K + (threadIdx.x % threads_per_row) * PACK_SIZE); + if constexpr (STAGES > 1) + { + uint32_t addr = cast_smem_ptr_to_uint(dst_ptr); + cp_async_cg_A(addr, src_ptr, local_mask & (ld_row + cta_offset_m < global_nrows)); + } + else + { + if (local_mask & (ld_row + cta_offset_m < global_nrows)) + *(uint4 *)dst_ptr = *src_ptr; + } + } +} + +template +__device__ __inline__ void global_to_share_one_stage_B(T *src, T *dst, int global_ncols, int cta_offset_m, int cta_offset_n, int cta_offset_k, int global_iter_k, int shared_iter_k, bool mask) +{ + constexpr int threads_needed = (CTA_N / kInterleave * CTA_K) / PACK_SIZE / SHARED_K_ITERS; + constexpr int threads_used = threads_needed < CTA_SIZE ? threads_needed : CTA_SIZE; + constexpr int total_global_iters = (CTA_N / kInterleave * CTA_K) / PACK_SIZE / threads_used; + constexpr int partial_global_iters = (total_global_iters + SHARED_K_ITERS - 1) / SHARED_K_ITERS; + constexpr int cta_step_m_or_n = (threads_used * PACK_SIZE) / CTA_K; + constexpr int warp_step_m_or_n = (WARP_SIZE * PACK_SIZE) / CTA_K; + constexpr int threads_per_row = CTA_K / PACK_SIZE; + constexpr int kSmemCol = CTA_K + SMEM_PAD_B; + bool local_mask = mask & (threadIdx.y * WARP_SIZE + threadIdx.x < threads_used); +#pragma unroll + for (int _global_iter = 0; _global_iter < partial_global_iters; ++_global_iter) + { + int global_iter = shared_iter_k * partial_global_iters + _global_iter; + + int ld_row = global_iter * cta_step_m_or_n + threadIdx.y * warp_step_m_or_n + (threadIdx.x / threads_per_row); + int ld_col = (threadIdx.x % threads_per_row); + int ld_col_swizzled = ld_col ^ (ld_row % 2) & 7; + void *dst_ptr = (void *)(dst + (ld_row * kSmemCol + ld_col_swizzled * PACK_SIZE)); + uint4 *src_ptr = (uint4 *)(src + global_iter_k * CTA_K + cta_offset_n / kInterleave * global_ncols + ld_row * global_ncols + ld_col * PACK_SIZE + cta_offset_k); + if constexpr (STAGES > 1) + { + uint32_t addr = cast_smem_ptr_to_uint(dst_ptr); + cp_async_cg_A(addr, src_ptr, local_mask); + } + else + { + if (local_mask) + *(uint4 *)dst_ptr = *src_ptr; + } + } +} + +template +__device__ __inline__ void global_to_share_one_stage_scales(T *src, T *dst, T *src_z, T *dst_z, int global_ncols, int cta_offset_m, int cta_offset_n, int cta_offset_k, int global_iter_k, int shared_iter_k, bool mask) +{ + constexpr int LD_AMOUNT = (G >= CTA_K) ? CTA_N : CTA_N * CTA_K / G; + constexpr int threads_needed = LD_AMOUNT / PACK_SIZE / 1; + constexpr int threads_used = threads_needed < CTA_SIZE ? threads_needed : CTA_SIZE; + constexpr int total_global_iters = LD_AMOUNT / PACK_SIZE / threads_used; + constexpr int threads_per_row = CTA_N / PACK_SIZE; + constexpr int kSmemCol = CTA_N; + bool local_mask = mask & (threadIdx.y * WARP_SIZE + threadIdx.x < threads_used); + int g_idx = (cta_offset_k + global_iter_k * CTA_K) / G; + + void *dst_ptr = (void *)(dst + (threadIdx.x / threads_per_row) * kSmemCol + (threadIdx.x % threads_per_row) * PACK_SIZE); + uint4 *src_ptr = (uint4 *)(src + g_idx * global_ncols + cta_offset_n + (threadIdx.x / threads_per_row) * global_ncols + (threadIdx.x % threads_per_row) * PACK_SIZE); + void *dst_ptr_z = (void *)(dst_z + (threadIdx.x / threads_per_row) * kSmemCol + (threadIdx.x % threads_per_row) * PACK_SIZE); + uint4 *src_ptr_z = (uint4 *)(src_z + g_idx * global_ncols + cta_offset_n + (threadIdx.x / threads_per_row) * global_ncols + (threadIdx.x % threads_per_row) * PACK_SIZE); + if (STAGES > 1) + { + uint32_t addr = cast_smem_ptr_to_uint(dst_ptr); + cp_async_cg_A(addr, src_ptr, local_mask); + uint32_t addr_z = cast_smem_ptr_to_uint(dst_ptr_z); + cp_async_cg_A(addr_z, src_ptr_z, local_mask); + } + else + { + if (local_mask) + { + *(uint4 *)dst_ptr = *src_ptr; + *(uint4 *)dst_ptr_z = *src_ptr_z; + } + } +} + +template +__device__ __inline__ void share_to_reg_one_stage_A(T *src, T *dst, int warp_offset_m, int warp_offset_n, int warp_offset_k, int k_0_1) +{ + constexpr int kSmemCol = CTA_K + SMEM_PAD_A; + + for (int shared_iter = 0; shared_iter < shared_iters; ++shared_iter) + { + + int ld_row = warp_offset_m + shared_iter * OP_M + (threadIdx.x % 16); + int ld_col = k_0_1 * 16 + (threadIdx.x / 16) * 8 + warp_offset_k; + int ld_col_swizzled = ((ld_col / PACK_SIZE) ^ (ld_row) & 7) * PACK_SIZE; + void *addr_ptr = (void *)(src + ld_row * kSmemCol + ld_col_swizzled); + + uint32_t addr = cast_smem_ptr_to_uint(addr_ptr); + ldmatrix_m8n8_x4_b16(dst, shared_iter, addr); + } +} + +template +__device__ __inline__ void share_to_reg_one_stage_B(T *src, T *src_scales, T *src_zeros, T *dst, T *dst_fp16, int warp_offset_m, int warp_offset_n, int warp_offset_k, int k_0_1) +{ + using T2 = typename std::conditional::value, half2, nv_bfloat162>::type; + constexpr int kSmemCol = CTA_K + SMEM_PAD_B; + int r0 = ((threadIdx.x / 8 / 2) * 8 + threadIdx.x % 8); + int c0 = ((threadIdx.x / 8) % 2) * 8; + int r = r0 / 4; + int c = (r0 % 4) * 16 + c0; + int c_swizzled = ((c / PACK_SIZE) ^ (r % 2) & 7) * PACK_SIZE; + + if constexpr (ldmatrix) + { +#pragma unroll + for (int shared_iter = 0; shared_iter < shared_iters; ++shared_iter) + { + void *addr_ptr = (void *)(src + warp_offset_n / kInterleave * kSmemCol + shared_iter * 16 / kInterleave * kSmemCol + k_0_1 * 16 + r * kSmemCol + c_swizzled + warp_offset_k); + uint32_t addr = cast_smem_ptr_to_uint(addr_ptr); + ldmatrix_m8n8_x4_b16(dst, shared_iter, addr); + } + } + +#pragma unroll + for (int shared_iter = 0; shared_iter < shared_iters; ++shared_iter) + { + T scale = src_scales[(warp_offset_k / G) * CTA_N + warp_offset_n + 16 * shared_iter + 8 * (k_0_1 % 2) + threadIdx.x / 4]; + T zero = src_zeros[(warp_offset_k / G) * CTA_N + warp_offset_n + 16 * shared_iter + 8 * (k_0_1 % 2) + threadIdx.x / 4]; + T2 scale2, zero2; + if constexpr (std::is_same::value) + { + scale2 = __half2half2(scale); + zero2 = __half2half2(zero); + } + else + { + scale2 = __bfloat162bfloat162(scale); + zero2 = __bfloat162bfloat162(zero); + } + T2 loaded[4]; + dequantize_s4_to_fp16x2(*reinterpret_cast(dst + (k_0_1 % 2) * 4 + (k_0_1 / 2 * 2) + shared_iter * 8), reinterpret_cast(loaded)); +#pragma unroll + for (int i = 0; i < 4; i++) + { + loaded[i] = __hfma2(loaded[i], scale2, zero2); + } + *reinterpret_cast(dst_fp16 + shared_iter * 16 + 8 * (k_0_1 % 2)) = *reinterpret_cast(loaded); + } +} + +template +__global__ void gemm_w4a16_T1(T *__restrict__ A, T *__restrict__ B, T *__restrict__ scales, T *__restrict__ zeros, T *__restrict__ C, int *__restrict__ semaphores, int M, int N, int K) +{ + using DTypeAccum = typename std::conditional::value, half, float>::type; + constexpr int NUM_WARPS_MN = CTA_M / WARP_M * CTA_N / WARP_N; + constexpr int NUM_WARPS = NUM_WARPS_MN * CTA_K / WARP_K; + constexpr int CTA_SIZE = NUM_WARPS * WARP_SIZE; + constexpr int CTA_SIZE_MN = NUM_WARPS_MN * WARP_SIZE; + constexpr int SLICES = CTA_K / WARP_K; + int num_blocks_n = (N + CTA_N - 1) / CTA_N; + int num_blocks_m = (M + CTA_M - 1) / CTA_M; + int blockIdx_x = 0; + int blockIdx_y = blockIdx.x % (num_blocks_m * num_blocks_n); + int blockIdx_z = blockIdx.x / (num_blocks_m * num_blocks_n); + const int log_tile = get_log_tile<1>((N + CTA_N - 1) / CTA_N); + int blockIdx_m = blockIdx_y / (num_blocks_n >> log_tile); + int blockIdx_n = blockIdx_y % (num_blocks_n >> log_tile); + const uint2 block_idx_mapping = get_block_idx_mapping(blockIdx_m, blockIdx_n, log_tile); + blockIdx_m = block_idx_mapping.x; + blockIdx_n = block_idx_mapping.y; + + DTypeAccum C_warp[CTA_M * CTA_N / CTA_SIZE_MN]; + constexpr int kSmemPadKA = CTA_K + SMEM_PAD_A; + constexpr int kSmemPadKB = CTA_K + SMEM_PAD_B; + constexpr int kSmemSizeAPerStage = CTA_M * kSmemPadKA; + constexpr int kSmemSizeBPerStage = CTA_N / kInterleave * kSmemPadKB; + constexpr int kSmemSizeA = kSmemSizeAPerStage * STAGES; + constexpr int kSmemSizeB = kSmemSizeBPerStage * STAGES; + constexpr int scales_load_interval = G >= CTA_K ? G / CTA_K : 1; + constexpr int scales_per_load = G < CTA_K ? CTA_K / G : 1; + constexpr int kSmemSizeScales = CTA_N * STAGES / scales_load_interval * scales_per_load; + constexpr int kSmemSizeZeros = CTA_N * STAGES / scales_load_interval * scales_per_load; + extern __shared__ half mem_shared[]; + T *A_shared = (T*)mem_shared; + T *B_shared = (T*)mem_shared + kSmemSizeA; + T *scales_shared = (T*)mem_shared + kSmemSizeA + kSmemSizeB; + T *zeros_shared = (T*)mem_shared + kSmemSizeA + kSmemSizeB + kSmemSizeScales; + T *C_shared = (T*)(mem_shared); + T A_shared_warp_[2][WARP_M * INTRIN_K / + WARP_SIZE]; + T B_shared_warp_[2][WARP_N * 32 / + WARP_SIZE]; + T B_shared_warp_tmp_[2][WARP_N * 16 / + WARP_SIZE]; + int cta_offset_m = blockIdx_m * CTA_M; + int cta_offset_n = blockIdx_n * CTA_N; + int cta_offset_k = blockIdx_z * (K / SPLITK); + int warp_mn = threadIdx.y % NUM_WARPS_MN; + int slice_id = threadIdx.y / NUM_WARPS_MN; + int warp_offset_n = (warp_mn % (CTA_N / WARP_N)) * WARP_N; + int warp_offset_m = (warp_mn / (CTA_N / WARP_N)) * WARP_M; + int warp_offset_k = slice_id * WARP_K; + + for (int i = 0; i < CTA_M * CTA_N / CTA_SIZE_MN; i++) + C_warp[i] = 0.0; + + int gemm_iters = (K + CTA_K - 1) / CTA_K / SPLITK; + int k_0_0_ld = 0; + int k_0_0 = 0; + constexpr int prologue_stages = STAGES == 1 ? 1 : STAGES - 1; +#pragma unroll + for (k_0_0_ld = 0; k_0_0_ld < prologue_stages; ++k_0_0_ld) + { + global_to_share_one_stage_A(A, A_shared + k_0_0_ld * kSmemSizeAPerStage, M, K, cta_offset_m, cta_offset_n, cta_offset_k, k_0_0_ld, 0, true); + global_to_share_one_stage_B(B, B_shared + k_0_0_ld * kSmemSizeBPerStage, K, cta_offset_m, cta_offset_n, cta_offset_k, k_0_0_ld, 0, true); + global_to_share_one_stage_scales( + scales, scales_shared + (k_0_0_ld / scales_load_interval * scales_per_load) * CTA_N, + zeros, zeros_shared + (k_0_0_ld / scales_load_interval * scales_per_load) * CTA_N, + N, cta_offset_m, cta_offset_n, cta_offset_k, + k_0_0_ld, 0, k_0_0_ld < gemm_iters && k_0_0_ld % scales_load_interval == 0); + if constexpr (STAGES > 1) + __pipeline_commit(); + } + if constexpr (STAGES > 1) + __pipeline_wait_prior(STAGES - 2); + __syncthreads(); + + share_to_reg_one_stage_A(A_shared, A_shared_warp_[0], warp_offset_m, warp_offset_n, warp_offset_k, 0); + share_to_reg_one_stage_B(B_shared, scales_shared, zeros_shared, B_shared_warp_tmp_[0], B_shared_warp_[0], warp_offset_m, warp_offset_n, warp_offset_k, 0); + constexpr int SHARED_K_ITERS = WARP_K / INTRIN_K; + + for (; k_0_0 < gemm_iters; ++k_0_0, ++k_0_0_ld) + { + int ld_stage = k_0_0_ld % STAGES; + int compute_stage = k_0_0 % STAGES; + T *A_shared_this_compute_stage; + T *B_shared_this_compute_stage; + T *scales_shared_this_compute_stage; + T *zeros_shared_this_compute_stage; + +#pragma unroll + for (int iter_k = 0; iter_k < SHARED_K_ITERS; ++iter_k) + { + A_shared_this_compute_stage = A_shared + compute_stage * kSmemSizeAPerStage; + B_shared_this_compute_stage = B_shared + compute_stage * kSmemSizeBPerStage; + scales_shared_this_compute_stage = scales_shared + (compute_stage / scales_load_interval * scales_per_load) * CTA_N; + zeros_shared_this_compute_stage = zeros_shared + (compute_stage / scales_load_interval * scales_per_load) * CTA_N; + share_to_reg_one_stage_A(A_shared_this_compute_stage, A_shared_warp_[(iter_k + 1) % 2], warp_offset_m, warp_offset_n, warp_offset_k, (iter_k + 1) % SHARED_K_ITERS); + if ((iter_k + 1) % kInterleave == 0) + { + if (compute_stage % 2 == 1) + { + share_to_reg_one_stage_B( + B_shared_this_compute_stage, scales_shared_this_compute_stage, zeros_shared_this_compute_stage, + B_shared_warp_tmp_[1], B_shared_warp_[((iter_k + 1) / 2) % 2], + warp_offset_m, warp_offset_n, warp_offset_k, (iter_k + 1) % SHARED_K_ITERS); + } + else + { + share_to_reg_one_stage_B( + B_shared_this_compute_stage, scales_shared_this_compute_stage, zeros_shared_this_compute_stage, + B_shared_warp_tmp_[0], B_shared_warp_[((iter_k + 1) / 2) % 2], + warp_offset_m, warp_offset_n, warp_offset_k, (iter_k + 1) % SHARED_K_ITERS); + } + } + else + { + if (compute_stage % 2 == 1) + { + share_to_reg_one_stage_B( + B_shared_this_compute_stage, scales_shared_this_compute_stage, zeros_shared_this_compute_stage, + B_shared_warp_tmp_[1], B_shared_warp_[((iter_k + 1) / 2) % 2], + warp_offset_m, warp_offset_n, warp_offset_k, (iter_k + 1) % SHARED_K_ITERS); + } + else + { + share_to_reg_one_stage_B( + B_shared_this_compute_stage, scales_shared_this_compute_stage, zeros_shared_this_compute_stage, + B_shared_warp_tmp_[0], B_shared_warp_[((iter_k + 1) / 2) % 2], + warp_offset_m, warp_offset_n, warp_offset_k, (iter_k + 1) % SHARED_K_ITERS); + } + } + T *A_shared_warp = A_shared_warp_[iter_k % 2]; + T *B_shared_warp = B_shared_warp_[(iter_k / 2) % 2]; + + for (int i_0_3 = 0; i_0_3 < WARP_M / INTRIN_M; ++i_0_3) + { + for (int j_0_4 = 0; j_0_4 < WARP_N / INTRIN_N; ++j_0_4) + { + if constexpr (std::is_same::value) + { + mma_m16n8k16_f16f16f16(C_warp + i_0_3 * WARP_N / INTRIN_N * 8 + j_0_4 * 8, A_shared_warp + i_0_3 * 8, B_shared_warp + j_0_4 * 16 + (iter_k % 2) * 4); + mma_m16n8k16_f16f16f16(C_warp + i_0_3 * WARP_N / INTRIN_N * 8 + j_0_4 * 8 + 4, A_shared_warp + i_0_3 * 8, B_shared_warp + j_0_4 * 16 + (iter_k % 2) * 4 + 8); + } + else + { + mma_m16n8k16_bf16bf16f32(C_warp + i_0_3 * WARP_N / INTRIN_N * 8 + j_0_4 * 8, A_shared_warp + i_0_3 * 8, B_shared_warp + j_0_4 * 16 + (iter_k % 2) * 4); + mma_m16n8k16_bf16bf16f32(C_warp + i_0_3 * WARP_N / INTRIN_N * 8 + j_0_4 * 8 + 4, A_shared_warp + i_0_3 * 8, B_shared_warp + j_0_4 * 16 + (iter_k % 2) * 4 + 8); + } + } + } + + if (iter_k < WARP_K / INTRIN_K - 1) + { + if constexpr (STAGES == 1) + __syncthreads(); + global_to_share_one_stage_A(A, A_shared + ld_stage * kSmemSizeAPerStage, M, K, cta_offset_m, cta_offset_n, cta_offset_k, k_0_0_ld, iter_k, k_0_0_ld < gemm_iters); + global_to_share_one_stage_B(B, B_shared + ld_stage * kSmemSizeBPerStage, K, cta_offset_m, cta_offset_n, cta_offset_k, k_0_0_ld, iter_k, k_0_0_ld < gemm_iters); + } + + if (iter_k == WARP_K / INTRIN_K - 2) + { + if constexpr (STAGES == 1 && WARP_K / INTRIN_K > 2) + { + __syncthreads(); + } + global_to_share_one_stage_A(A, A_shared + ld_stage * kSmemSizeAPerStage, M, K, cta_offset_m, cta_offset_n, cta_offset_k, k_0_0_ld, iter_k + 1, k_0_0_ld < gemm_iters); + global_to_share_one_stage_B(B, B_shared + ld_stage * kSmemSizeBPerStage, K, cta_offset_m, cta_offset_n, cta_offset_k, k_0_0_ld, iter_k + 1, k_0_0_ld < gemm_iters); + global_to_share_one_stage_scales( + scales, scales_shared + (ld_stage / scales_load_interval * scales_per_load) * CTA_N, + zeros, zeros_shared + (ld_stage / scales_load_interval * scales_per_load) * CTA_N, + N, cta_offset_m, cta_offset_n, cta_offset_k, + k_0_0_ld, iter_k, k_0_0_ld < gemm_iters && k_0_0_ld % scales_load_interval == 0); + if constexpr (STAGES > 1) + { + __pipeline_commit(); + __pipeline_wait_prior(STAGES - 2); + } + compute_stage = (k_0_0 + 1) % STAGES; + __syncthreads(); + } + } + } + __pipeline_commit(); + __pipeline_wait_prior(0); + __syncthreads(); + + if constexpr (std::is_same::value) + { + if constexpr (SLICES > 1) + { + #pragma unroll + for (int z = 0; z < SLICES; ++z) + { + if (slice_id == z) + { + #pragma unroll + for (int ax0_0_1 = 0; ax0_0_1 < WARP_M / INTRIN_M; ++ax0_0_1) + { + #pragma unroll + for (int ax1_0_1 = 0; ax1_0_1 < WARP_N / INTRIN_N; ++ax1_0_1) + { + #pragma unroll + for (int local_id = 0; local_id < OP_M * 16 / WARP_SIZE; ++local_id) + { + if (z > 0) + { + C_warp[ax0_0_1 * WARP_N / INTRIN_N * 8 + ax1_0_1 * 8 + local_id] += C_shared[warp_offset_m * CTA_N + ax0_0_1 * OP_M * CTA_N + warp_offset_n + ax1_0_1 * 16 + ((local_id % 4) / 2 * 8 + (threadIdx.x / 4)) * CTA_N + (local_id / 4) * 8 + (local_id % 2) + (threadIdx.x % 4) * 2]; + } + C_shared[warp_offset_m * CTA_N + ax0_0_1 * OP_M * CTA_N + warp_offset_n + ax1_0_1 * 16 + ((local_id % 4) / 2 * 8 + (threadIdx.x / 4)) * CTA_N + (local_id / 4) * 8 + (local_id % 2) + (threadIdx.x % 4) * 2] = C_warp[ax0_0_1 * WARP_N / INTRIN_N * 8 + ax1_0_1 * 8 + local_id]; + }; + } + } + } + __syncthreads(); + } + if (slice_id == 0) + { + #pragma unroll + for (int ax0_0_1 = 0; ax0_0_1 < WARP_M / INTRIN_M; ++ax0_0_1) + { + #pragma unroll + for (int ax1_0_1 = 0; ax1_0_1 < WARP_N / INTRIN_N; ++ax1_0_1) + { + #pragma unroll + for (int local_id = 0; local_id < OP_M * 16 / WARP_SIZE; ++local_id) + { + C_warp[ax0_0_1 * WARP_N / INTRIN_N * 8 + ax1_0_1 * 8 + local_id] = C_shared[warp_offset_m * CTA_N + ax0_0_1 * OP_M * CTA_N + warp_offset_n + ax1_0_1 * 16 + ((local_id % 4) / 2 * 8 + (threadIdx.x / 4)) * CTA_N + (local_id / 4) * 8 + (local_id % 2) + (threadIdx.x % 4) * 2]; + }; + } + } + } + } + + if (slice_id == 0) + { + Semaphore semaphore(semaphores + blockIdx_y, threadIdx.x); + + if constexpr (SPLITK > 1) + { + semaphore.fetch(); + } + + if (blockIdx_z != 0) + { + semaphore.wait(blockIdx_z); + for (int ax0_0_1 = 0; ax0_0_1 < WARP_M / INTRIN_M; ++ax0_0_1) + { + for (int ax1_0_1 = 0; ax1_0_1 < WARP_N / INTRIN_N; ++ax1_0_1) + { + for (int local_id = 0; local_id < OP_M * 16 / WARP_SIZE; local_id += 2) + { + int write_row = cta_offset_m + warp_offset_m + ax0_0_1 * OP_M + ((local_id % 4) / 2 * 8 + (threadIdx.x / 4)); + + if (write_row < M) + { + half2 *existing_psum_ptr = reinterpret_cast( + C + write_row * N + + cta_offset_n + warp_offset_n + ax1_0_1 * 16 + + (local_id / 4) * 8 + (local_id % 2) + (threadIdx.x % 4) * 2); + + *existing_psum_ptr = __hadd2(*existing_psum_ptr, + *reinterpret_cast(C_warp + ax0_0_1 * WARP_N / INTRIN_N * 8 + + ax1_0_1 * 8 + local_id)); + } + }; + } + } + } + else + { + for (int ax0_0_1 = 0; ax0_0_1 < WARP_M / INTRIN_M; ++ax0_0_1) + { + for (int ax1_0_1 = 0; ax1_0_1 < WARP_N / INTRIN_N; ++ax1_0_1) + { + for (int local_id = 0; local_id < OP_M * 16 / WARP_SIZE; local_id += 2) + { + int write_row = cta_offset_m + warp_offset_m + ax0_0_1 * OP_M + ((local_id % 4) / 2 * 8 + (threadIdx.x / 4)); + if (write_row < M) + { + *reinterpret_cast( + C + write_row * N + + cta_offset_n + warp_offset_n + ax1_0_1 * 16 + + (local_id / 4) * 8 + (local_id % 2) + (threadIdx.x % 4) * 2) = + *reinterpret_cast(C_warp + ax0_0_1 * WARP_N / INTRIN_N * 8 + + ax1_0_1 * 8 + local_id); + } + }; + } + } + } + + if constexpr (SPLITK > 1) + { + + int lock = 0; + if (SPLITK == blockIdx_z + 1) + { + + lock = 0; + } + else + { + lock = blockIdx_z + 1; + } + semaphore.release(lock); + } + } + } + else + { + // first convert fp32 to bf16 + nv_bfloat16 C_warp16[CTA_M * CTA_N / CTA_SIZE_MN]; +#pragma unroll + for (int i = 0; i < CTA_M * CTA_N / CTA_SIZE_MN / 2; ++i) + { + ((nv_bfloat162*)C_warp16)[i] = __float22bfloat162_rn(((float2*)C_warp)[i]); + } + + // the following is the same as fp16. Maybe there is a neat way to implement this. + if constexpr (SLICES > 1) + { +#pragma unroll + for (int z = 0; z < SLICES; ++z) + { + if (slice_id == z) + { +#pragma unroll + for (int ax0_0_1 = 0; ax0_0_1 < WARP_M / INTRIN_M; ++ax0_0_1) + { +#pragma unroll + for (int ax1_0_1 = 0; ax1_0_1 < WARP_N / INTRIN_N; ++ax1_0_1) + { +#pragma unroll + for (int local_id = 0; local_id < OP_M * 16 / WARP_SIZE; ++local_id) + { + if (z > 0) + { + C_warp16[ax0_0_1 * WARP_N / INTRIN_N * 8 + ax1_0_1 * 8 + local_id] += C_shared[warp_offset_m * CTA_N + ax0_0_1 * OP_M * CTA_N + warp_offset_n + ax1_0_1 * 16 + ((local_id % 4) / 2 * 8 + (threadIdx.x / 4)) * CTA_N + (local_id / 4) * 8 + (local_id % 2) + (threadIdx.x % 4) * 2]; + } + C_shared[warp_offset_m * CTA_N + ax0_0_1 * OP_M * CTA_N + warp_offset_n + ax1_0_1 * 16 + ((local_id % 4) / 2 * 8 + (threadIdx.x / 4)) * CTA_N + (local_id / 4) * 8 + (local_id % 2) + (threadIdx.x % 4) * 2] = C_warp16[ax0_0_1 * WARP_N / INTRIN_N * 8 + ax1_0_1 * 8 + local_id]; + }; + } + } + } + __syncthreads(); + } + if (slice_id == 0) + { +#pragma unroll + for (int ax0_0_1 = 0; ax0_0_1 < WARP_M / INTRIN_M; ++ax0_0_1) + { +#pragma unroll + for (int ax1_0_1 = 0; ax1_0_1 < WARP_N / INTRIN_N; ++ax1_0_1) + { +#pragma unroll + for (int local_id = 0; local_id < OP_M * 16 / WARP_SIZE; ++local_id) + { + C_warp16[ax0_0_1 * WARP_N / INTRIN_N * 8 + ax1_0_1 * 8 + local_id] = C_shared[warp_offset_m * CTA_N + ax0_0_1 * OP_M * CTA_N + warp_offset_n + ax1_0_1 * 16 + ((local_id % 4) / 2 * 8 + (threadIdx.x / 4)) * CTA_N + (local_id / 4) * 8 + (local_id % 2) + (threadIdx.x % 4) * 2]; + }; + } + } + } + } + + if (slice_id == 0) + { + Semaphore semaphore(semaphores + blockIdx_y, threadIdx.x); + + if constexpr (SPLITK > 1) + { + semaphore.fetch(); + } + + if (blockIdx_z != 0) + { + semaphore.wait(blockIdx_z); + for (int ax0_0_1 = 0; ax0_0_1 < WARP_M / INTRIN_M; ++ax0_0_1) + { + for (int ax1_0_1 = 0; ax1_0_1 < WARP_N / INTRIN_N; ++ax1_0_1) + { + for (int local_id = 0; local_id < OP_M * 16 / WARP_SIZE; local_id += 2) + { + int write_row = cta_offset_m + warp_offset_m + ax0_0_1 * OP_M + ((local_id % 4) / 2 * 8 + (threadIdx.x / 4)); + + if (write_row < M) + { + nv_bfloat162 *existing_psum_ptr = reinterpret_cast( + C + write_row * N + + cta_offset_n + warp_offset_n + ax1_0_1 * 16 + + (local_id / 4) * 8 + (local_id % 2) + (threadIdx.x % 4) * 2); + + *existing_psum_ptr = __hadd2(*existing_psum_ptr, + *reinterpret_cast(C_warp16 + ax0_0_1 * WARP_N / INTRIN_N * 8 + + ax1_0_1 * 8 + local_id)); + } + }; + } + } + } + else + { + for (int ax0_0_1 = 0; ax0_0_1 < WARP_M / INTRIN_M; ++ax0_0_1) + { + for (int ax1_0_1 = 0; ax1_0_1 < WARP_N / INTRIN_N; ++ax1_0_1) + { + for (int local_id = 0; local_id < OP_M * 16 / WARP_SIZE; local_id += 2) + { + int write_row = cta_offset_m + warp_offset_m + ax0_0_1 * OP_M + ((local_id % 4) / 2 * 8 + (threadIdx.x / 4)); + if (write_row < M) + { + *reinterpret_cast( + C + write_row * N + + cta_offset_n + warp_offset_n + ax1_0_1 * 16 + + (local_id / 4) * 8 + (local_id % 2) + (threadIdx.x % 4) * 2) = + *reinterpret_cast(C_warp16 + ax0_0_1 * WARP_N / INTRIN_N * 8 + + ax1_0_1 * 8 + local_id); + } + }; + } + } + } + + if constexpr (SPLITK > 1) + { + + int lock = 0; + if (SPLITK == blockIdx_z + 1) + { + + lock = 0; + } + else + { + lock = blockIdx_z + 1; + } + semaphore.release(lock); + } + } + } +} + +template +__device__ __inline__ void global_to_share_one_stage_A_T2(T *src, T *dst, int global_nrows, int global_ncols, int cta_offset_m, int cta_offset_n, int global_iter_k, int shared_iter_k, bool mask) +{ + constexpr int threads_needed = (CTA_M * CTA_K) / PACK_SIZE / SHARED_K_ITERS; + constexpr int threads_used = threads_needed < CTA_SIZE ? threads_needed : CTA_SIZE; + constexpr int total_global_iters = (CTA_M * CTA_K) / PACK_SIZE / threads_used; + constexpr int partial_global_iters = (total_global_iters + SHARED_K_ITERS - 1) / SHARED_K_ITERS; + constexpr int cta_step_m_or_n = (threads_used * PACK_SIZE) / CTA_K; + constexpr int warp_step_m_or_n = (WARP_SIZE * PACK_SIZE) / CTA_K; + constexpr int threads_per_row = CTA_K / PACK_SIZE; + constexpr int kSmemCol = CTA_K + SMEM_PAD_A; + bool local_mask = mask & (threadIdx.y * WARP_SIZE + threadIdx.x < threads_used); + int ld_col = (threadIdx.x % threads_per_row); +#pragma unroll + for (int _global_iter = 0; _global_iter < partial_global_iters; ++_global_iter) + { + int global_iter = shared_iter_k * partial_global_iters + _global_iter; + int ld_row = global_iter * cta_step_m_or_n + threadIdx.y * warp_step_m_or_n + (threadIdx.x / threads_per_row); + int ld_col_swizzled = (ld_col ^ (ld_row) & 7) * PACK_SIZE; + void *dst_ptr = (void *)(dst + ld_row * kSmemCol + ld_col_swizzled); + uint4 *src_ptr = (uint4 *)(src + (ld_row + cta_offset_m) * global_ncols + ld_col * PACK_SIZE + global_iter_k * CTA_K); // cta_offset_m * global_ncols + global_iter * cta_step_m_or_n * global_ncols + threadIdx.y * warp_step_m_or_n * global_ncols + (threadIdx.x / threads_per_row) * global_ncols + global_iter_k * CTA_K + (threadIdx.x % threads_per_row) * PACK_SIZE); + if constexpr (STAGES > 1) + { + uint32_t addr = cast_smem_ptr_to_uint(dst_ptr); + cp_async_cg_A(addr, src_ptr, local_mask & (ld_row + cta_offset_m < global_nrows)); + } + else + { + if (local_mask & (ld_row + cta_offset_m < global_nrows)) + *(uint4 *)dst_ptr = *src_ptr; + } + } +} + +template +__device__ __inline__ void global_to_share_one_stage_B_T2(T *src, T *dst, int global_ncols, int cta_offset_m, int cta_offset_n, int global_iter_k, int shared_iter_k, bool mask) +{ + constexpr int threads_needed = (CTA_N / kInterleave * CTA_K) / PACK_SIZE / SHARED_K_ITERS; + constexpr int threads_used = threads_needed < CTA_SIZE ? threads_needed : CTA_SIZE; + constexpr int total_global_iters = (CTA_N / kInterleave * CTA_K) / PACK_SIZE / threads_used; + constexpr int partial_global_iters = (total_global_iters + SHARED_K_ITERS - 1) / SHARED_K_ITERS; + constexpr int cta_step_m_or_n = (threads_used * PACK_SIZE) / CTA_K; + constexpr int warp_step_m_or_n = (WARP_SIZE * PACK_SIZE) / CTA_K; + constexpr int threads_per_row = CTA_K / PACK_SIZE; + constexpr int kSmemCol = CTA_K + SMEM_PAD_B; + bool local_mask = mask & (threadIdx.y * WARP_SIZE + threadIdx.x < threads_used); +#pragma unroll + for (int _global_iter = 0; _global_iter < partial_global_iters; ++_global_iter) + { + int global_iter = shared_iter_k * partial_global_iters + _global_iter; + + int ld_row = global_iter * cta_step_m_or_n + threadIdx.y * warp_step_m_or_n + (threadIdx.x / threads_per_row); + int ld_col = (threadIdx.x % threads_per_row); + int ld_col_swizzled = ld_col ^ (ld_row % 2) & 7; + void *dst_ptr = (void *)(dst + (ld_row * kSmemCol + ld_col_swizzled * PACK_SIZE)); + uint4 *src_ptr = (uint4 *)(src + global_iter_k * CTA_K + cta_offset_n / kInterleave * global_ncols + ld_row * global_ncols + ld_col * PACK_SIZE); + if constexpr (STAGES > 1) + { + uint32_t addr = cast_smem_ptr_to_uint(dst_ptr); + cp_async_cg_A(addr, src_ptr, local_mask); + } + else + { + if (local_mask) + *(uint4 *)dst_ptr = *src_ptr; + } + } +} + +template +__device__ __inline__ void global_to_share_one_stage_scales_T2(T *src, T *dst, T *src_z, T *dst_z, int global_ncols, int cta_offset_m, int cta_offset_n, int global_iter_k, int shared_iter_k, bool mask) +{ + constexpr int threads_needed = CTA_N / PACK_SIZE / 1; + constexpr int threads_used = threads_needed < CTA_SIZE ? threads_needed : CTA_SIZE; + constexpr int total_global_iters = CTA_N / PACK_SIZE / threads_used; + constexpr int threads_per_row = CTA_N / PACK_SIZE; + constexpr int kSmemCol = CTA_N; + bool local_mask = mask & (threadIdx.y * WARP_SIZE + threadIdx.x < threads_used); + int g_idx = global_iter_k * CTA_K / G; + + void *dst_ptr = (void *)(dst + (threadIdx.x % threads_per_row) * PACK_SIZE); + uint4 *src_ptr = (uint4 *)(src + g_idx * global_ncols + cta_offset_n + (threadIdx.x % threads_per_row) * PACK_SIZE); + void *dst_ptr_z = (void *)(dst_z + (threadIdx.x % threads_per_row) * PACK_SIZE); + uint4 *src_ptr_z = (uint4 *)(src_z + g_idx * global_ncols + cta_offset_n + (threadIdx.x % threads_per_row) * PACK_SIZE); + if (STAGES > 1) + { + uint32_t addr = cast_smem_ptr_to_uint(dst_ptr); + cp_async_cg_A(addr, src_ptr, local_mask); + uint32_t addr_z = cast_smem_ptr_to_uint(dst_ptr_z); + cp_async_cg_A(addr_z, src_ptr_z, local_mask); + } + else + { + if (local_mask) + { + *(uint4 *)dst_ptr = *src_ptr; + *(uint4 *)dst_ptr_z = *src_ptr_z; + } + } +} + +template +__device__ __inline__ void share_to_reg_one_stage_A_T2(T *src, T *dst, int warp_offset_m, int warp_offset_n, int k_0_1) +{ + constexpr int kSmemCol = CTA_K + SMEM_PAD_A; + + for (int shared_iter = 0; shared_iter < shared_iters; ++shared_iter) + { + + int ld_row = warp_offset_m + shared_iter * OP_M + (threadIdx.x % 16); + int ld_col = k_0_1 * 16 + (threadIdx.x / 16) * 8; + int ld_col_swizzled = ((ld_col / PACK_SIZE) ^ (ld_row) & 7) * PACK_SIZE; + void *addr_ptr = (void *)(src + ld_row * kSmemCol + ld_col_swizzled); + + uint32_t addr = cast_smem_ptr_to_uint(addr_ptr); + ldmatrix_m8n8_x4_b16(dst, shared_iter, addr); + } +} + +template +__device__ __inline__ void share_to_reg_one_stage_B_T2(T *src, T *src_scales, T *src_zeros, T *dst, T *dst_fp16, int warp_offset_m, int warp_offset_n, int k_0_1) +{ + using T2 = typename std::conditional::value, half2, nv_bfloat162>::type; + constexpr int kSmemCol = CTA_K + SMEM_PAD_B; + int r0 = ((threadIdx.x / 8 / 2) * 8 + threadIdx.x % 8); + int c0 = ((threadIdx.x / 8) % 2) * 8; + int r = r0 / 4; + int c = (r0 % 4) * 16 + c0; + int c_swizzled = ((c / PACK_SIZE) ^ (r % 2) & 7) * PACK_SIZE; + + if constexpr (ldmatrix) + { +#pragma unroll + for (int shared_iter = 0; shared_iter < shared_iters; ++shared_iter) + { + void *addr_ptr = (void *)(src + warp_offset_n / kInterleave * kSmemCol + shared_iter * 16 / kInterleave * kSmemCol + k_0_1 * 16 + r * kSmemCol + c_swizzled); + uint32_t addr = cast_smem_ptr_to_uint(addr_ptr); + ldmatrix_m8n8_x4_b16(dst, shared_iter, addr); + } + } + +#pragma unroll + for (int shared_iter = 0; shared_iter < shared_iters; ++shared_iter) + { + T scale = src_scales[warp_offset_n + 16 * shared_iter + 8 * (k_0_1 % 2) + threadIdx.x / 4]; + T zero = src_zeros[warp_offset_n + 16 * shared_iter + 8 * (k_0_1 % 2) + threadIdx.x / 4]; + T2 scale2, zero2; + if constexpr (std::is_same::value) + { + scale2 = __half2half2(scale); + zero2 = __half2half2(zero); + } + else + { + scale2 = __bfloat162bfloat162(scale); + zero2 = __bfloat162bfloat162(zero); + } + T2 loaded[4]; + dequantize_s4_to_fp16x2(*reinterpret_cast(dst + (k_0_1 % 2) * 4 + (k_0_1 / 2 * 2) + shared_iter * 8), reinterpret_cast(loaded)); +#pragma unroll + for (int i = 0; i < 4; i++) + { + loaded[i] = __hfma2(loaded[i], scale2, zero2); + } + *reinterpret_cast(dst_fp16 + shared_iter * 16 + 8 * (k_0_1 % 2)) = *reinterpret_cast(loaded); + } +} + +template +__global__ void gemm_w4a16_T2(T *__restrict__ A, T *__restrict__ B, T *__restrict__ scales, T *__restrict__ zeros, T *__restrict__ C, int M, int N, int K) +{ + using DTypeAccum = typename std::conditional::value, half, float>::type; + constexpr int NUM_WARPS = CTA_M / WARP_M * CTA_N / WARP_N; + constexpr int CTA_SIZE = NUM_WARPS * WARP_SIZE; + int num_blocks_n = (N + CTA_N - 1) / CTA_N; + int num_blocks_m = (M + CTA_M - 1) / CTA_M; + int blockIdx_x = 0; + int blockIdx_y = blockIdx.x % (num_blocks_m * num_blocks_n); + int blockIdx_z = blockIdx.x / (num_blocks_m * num_blocks_n); + const int log_tile = get_log_tile<1>((N + CTA_N - 1) / CTA_N); + int blockIdx_m = blockIdx_y / (num_blocks_n >> log_tile); + int blockIdx_n = blockIdx_y % (num_blocks_n >> log_tile); + const uint2 block_idx_mapping = get_block_idx_mapping(blockIdx_m, blockIdx_n, log_tile); + blockIdx_m = block_idx_mapping.x; + blockIdx_n = block_idx_mapping.y; + + DTypeAccum C_warp[CTA_M * CTA_N / CTA_SIZE]; + constexpr int kSmemPadKA = CTA_K + SMEM_PAD_A; + constexpr int kSmemPadKB = CTA_K + SMEM_PAD_B; + constexpr int kSmemSizeAPerStage = CTA_M * kSmemPadKA; + constexpr int kSmemSizeBPerStage = CTA_N / kInterleave * kSmemPadKB; + constexpr int kSmemSizeA = kSmemSizeAPerStage * STAGES; + constexpr int kSmemSizeB = kSmemSizeBPerStage * STAGES; + constexpr int kSmemSizeScales = CTA_N * STAGES / 2; + constexpr int kSmemSizeZeros = CTA_N * STAGES / 2; + constexpr int scales_load_interval = G / CTA_K; + extern __shared__ half mem_shared[]; + T *A_shared = (T*)mem_shared; + T *B_shared = (T*)mem_shared + kSmemSizeA; + T *scales_shared = (T*)mem_shared + kSmemSizeA + kSmemSizeB; + T *zeros_shared = (T*)mem_shared + kSmemSizeA + kSmemSizeB + kSmemSizeScales; + T A_shared_warp_[2][WARP_M * INTRIN_K / + WARP_SIZE]; + T B_shared_warp_[2][WARP_N * 32 / + WARP_SIZE]; + T B_shared_warp_tmp_[2][WARP_N * 16 / + WARP_SIZE]; + int cta_offset_m = blockIdx_m * CTA_M; + int cta_offset_n = blockIdx_n * CTA_N; + int warp_offset_m = (threadIdx.y % (CTA_M / WARP_M)) * WARP_M; + int warp_offset_n = (threadIdx.y / (CTA_M / WARP_M)) * WARP_N; + + for (int i = 0; i < CTA_M * CTA_N / CTA_SIZE; i++) + C_warp[i] = 0.0; + + int gemm_iters = (K + CTA_K - 1) / CTA_K; + int k_0_0_ld = 0; + int k_0_0 = 0; + constexpr int prologue_stages = STAGES == 1 ? 1 : STAGES - 1; +#pragma unroll + for (k_0_0_ld = 0; k_0_0_ld < prologue_stages; ++k_0_0_ld) + { + global_to_share_one_stage_A_T2(A, A_shared + k_0_0_ld * kSmemSizeAPerStage, M, K, cta_offset_m, cta_offset_n, k_0_0_ld, 0, true); + global_to_share_one_stage_B_T2(B, B_shared + k_0_0_ld * kSmemSizeBPerStage, K, cta_offset_m, cta_offset_n, k_0_0_ld, 0, true); + global_to_share_one_stage_scales_T2( + scales, scales_shared + (k_0_0_ld / scales_load_interval) * CTA_N, + zeros, zeros_shared + (k_0_0_ld / scales_load_interval) * CTA_N, + N, cta_offset_m, cta_offset_n, k_0_0_ld, 0, k_0_0_ld < gemm_iters && k_0_0_ld % scales_load_interval == 0); + if constexpr (STAGES > 1) + __pipeline_commit(); + } + if constexpr (STAGES > 1) + __pipeline_wait_prior(STAGES - 2); + __syncthreads(); + + share_to_reg_one_stage_A_T2(A_shared, A_shared_warp_[0], warp_offset_m, warp_offset_n, 0); + share_to_reg_one_stage_B_T2(B_shared, scales_shared, zeros_shared, B_shared_warp_tmp_[0], B_shared_warp_[0], warp_offset_m, warp_offset_n, 0); + constexpr int SHARED_K_ITERS = WARP_K / INTRIN_K; + + for (; k_0_0 < gemm_iters; ++k_0_0, ++k_0_0_ld) + { + int ld_stage = k_0_0_ld % STAGES; + int compute_stage = k_0_0 % STAGES; + T *A_shared_this_compute_stage; + T *B_shared_this_compute_stage; + T *scales_shared_this_compute_stage; + T *zeros_shared_this_compute_stage; + + for (int iter_k = 0; iter_k < SHARED_K_ITERS; ++iter_k) + { + A_shared_this_compute_stage = A_shared + compute_stage * kSmemSizeAPerStage; + B_shared_this_compute_stage = B_shared + compute_stage * kSmemSizeBPerStage; + scales_shared_this_compute_stage = scales_shared + (compute_stage / scales_load_interval) * CTA_N; + zeros_shared_this_compute_stage = zeros_shared + (compute_stage / scales_load_interval) * CTA_N; + share_to_reg_one_stage_A_T2(A_shared_this_compute_stage, A_shared_warp_[(iter_k + 1) % 2], warp_offset_m, warp_offset_n, (iter_k + 1) % SHARED_K_ITERS); + if ((iter_k + 1) % kInterleave == 0) + { + if (compute_stage % 2 == 1) + { + share_to_reg_one_stage_B_T2( + B_shared_this_compute_stage, scales_shared_this_compute_stage, zeros_shared_this_compute_stage, + B_shared_warp_tmp_[1], B_shared_warp_[((iter_k + 1) / 2) % 2], + warp_offset_m, warp_offset_n, (iter_k + 1) % SHARED_K_ITERS); + } + else + { + share_to_reg_one_stage_B_T2( + B_shared_this_compute_stage, scales_shared_this_compute_stage, zeros_shared_this_compute_stage, + B_shared_warp_tmp_[0], B_shared_warp_[((iter_k + 1) / 2) % 2], + warp_offset_m, warp_offset_n, (iter_k + 1) % SHARED_K_ITERS); + } + } + else + { + if (compute_stage % 2 == 1) + { + share_to_reg_one_stage_B_T2( + B_shared_this_compute_stage, scales_shared_this_compute_stage, zeros_shared_this_compute_stage, + B_shared_warp_tmp_[1], B_shared_warp_[((iter_k + 1) / 2) % 2], + warp_offset_m, warp_offset_n, (iter_k + 1) % SHARED_K_ITERS); + } + else + { + share_to_reg_one_stage_B_T2( + B_shared_this_compute_stage, scales_shared_this_compute_stage, zeros_shared_this_compute_stage, + B_shared_warp_tmp_[0], B_shared_warp_[((iter_k + 1) / 2) % 2], + warp_offset_m, warp_offset_n, (iter_k + 1) % SHARED_K_ITERS); + } + } + __syncthreads(); + T *A_shared_warp = A_shared_warp_[iter_k % 2]; + T *B_shared_warp = B_shared_warp_[(iter_k / 2) % 2]; + for (int i_0_3 = 0; i_0_3 < WARP_M / INTRIN_M; ++i_0_3) + { + for (int j_0_4 = 0; j_0_4 < WARP_N / INTRIN_N; ++j_0_4) + { + if constexpr (std::is_same::value) + { + mma_m16n8k16_f16f16f16(C_warp + i_0_3 * WARP_N / INTRIN_N * 8 + j_0_4 * 8, A_shared_warp + i_0_3 * 8, B_shared_warp + j_0_4 * 16 + (iter_k % 2) * 4); + mma_m16n8k16_f16f16f16(C_warp + i_0_3 * WARP_N / INTRIN_N * 8 + j_0_4 * 8 + 4, A_shared_warp + i_0_3 * 8, B_shared_warp + j_0_4 * 16 + (iter_k % 2) * 4 + 8); + } + else + { + mma_m16n8k16_bf16bf16f32(C_warp + i_0_3 * WARP_N / INTRIN_N * 8 + j_0_4 * 8, A_shared_warp + i_0_3 * 8, B_shared_warp + j_0_4 * 16 + (iter_k % 2) * 4); + mma_m16n8k16_bf16bf16f32(C_warp + i_0_3 * WARP_N / INTRIN_N * 8 + j_0_4 * 8 + 4, A_shared_warp + i_0_3 * 8, B_shared_warp + j_0_4 * 16 + (iter_k % 2) * 4 + 8); + } + } + } + + if (iter_k < WARP_K / INTRIN_K - 1) + { + if constexpr (STAGES == 1) + __syncthreads(); + global_to_share_one_stage_A_T2(A, A_shared + ld_stage * kSmemSizeAPerStage, M, K, cta_offset_m, cta_offset_n, k_0_0_ld, iter_k, k_0_0_ld < gemm_iters); + global_to_share_one_stage_B_T2(B, B_shared + ld_stage * kSmemSizeBPerStage, K, cta_offset_m, cta_offset_n, k_0_0_ld, iter_k, k_0_0_ld < gemm_iters); + } + + if (iter_k == WARP_K / INTRIN_K - 2) + { + if constexpr (STAGES == 1 && WARP_K / INTRIN_K > 2) + { + __syncthreads(); + } + global_to_share_one_stage_A_T2(A, A_shared + ld_stage * kSmemSizeAPerStage, M, K, cta_offset_m, cta_offset_n, k_0_0_ld, iter_k + 1, k_0_0_ld < gemm_iters); + global_to_share_one_stage_B_T2(B, B_shared + ld_stage * kSmemSizeBPerStage, K, cta_offset_m, cta_offset_n, k_0_0_ld, iter_k + 1, k_0_0_ld < gemm_iters); + global_to_share_one_stage_scales_T2( + scales, scales_shared + (ld_stage / scales_load_interval) * CTA_N, + zeros, zeros_shared + (ld_stage / scales_load_interval) * CTA_N, + N, cta_offset_m, cta_offset_n, k_0_0_ld, iter_k, k_0_0_ld < gemm_iters && k_0_0_ld % scales_load_interval == 0); + if constexpr (STAGES > 1) + { + __pipeline_commit(); + __pipeline_wait_prior(STAGES - 2); + } + compute_stage = (k_0_0 + 1) % STAGES; + __syncthreads(); + } + } + } + for (int ax0_0_1 = 0; ax0_0_1 < WARP_M / INTRIN_M; ++ax0_0_1) + { + for (int ax1_0_1 = 0; ax1_0_1 < WARP_N / INTRIN_N; ++ax1_0_1) + { + for (int local_id = 0; local_id < OP_M * 16 / WARP_SIZE; local_id += 2) + { + int write_row = cta_offset_m + warp_offset_m + ax0_0_1 * OP_M + ((local_id % 4) / 2 * 8 + (threadIdx.x / 4)); + if (write_row < M) + { + if constexpr (std::is_same::value) + { + *reinterpret_cast( + C + write_row * N + + cta_offset_n + warp_offset_n + ax1_0_1 * 16 + + (local_id / 4) * 8 + (local_id % 2) + (threadIdx.x % 4) * 2) = + (*reinterpret_cast(C_warp + ax0_0_1 * WARP_N / INTRIN_N * 8 + + ax1_0_1 * 8 + local_id)); + } + else + { + *reinterpret_cast( + C + write_row * N + + cta_offset_n + warp_offset_n + ax1_0_1 * 16 + + (local_id / 4) * 8 + (local_id % 2) + (threadIdx.x % 4) * 2) = + (__float22bfloat162_rn(*reinterpret_cast(C_warp + ax0_0_1 * WARP_N / INTRIN_N * 8 + + ax1_0_1 * 8 + local_id))); + } + } + }; + } + } +} + +torch::Tensor gemm_forward_cuda_new( + torch::Tensor _in_feats, + torch::Tensor _kernel, + torch::Tensor _scales, + torch::Tensor _zeros) +{ + std::vector output_shape = _in_feats.sizes().vec(); + output_shape.back() = _kernel.size(0) * kInterleave; + int num_in_feats = _in_feats.numel() / _in_feats.size(-1); + int num_in_channels = _in_feats.size(-1); + auto options = + torch::TensorOptions().dtype(_in_feats.dtype()).device(_in_feats.device()); + auto options_int = + torch::TensorOptions().dtype(torch::kInt32).device(_in_feats.device()); + at::Tensor _out_feats = torch::empty(output_shape, options); + int num_out_feats = _out_feats.numel() / _out_feats.size(-1); + int num_out_channels = _out_feats.size(-1); + + auto data_type = _in_feats.scalar_type(); + TORCH_CHECK(_scales.scalar_type() == data_type); + TORCH_CHECK(_zeros.scalar_type() == data_type); + + DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP16(data_type, ctype, { + auto in_feats = reinterpret_cast(_in_feats.data_ptr()); + auto kernel = reinterpret_cast(_kernel.data_ptr()); + auto scales = reinterpret_cast(_scales.data_ptr()); + auto zeros = reinterpret_cast(_zeros.data_ptr()); + auto out_feats = reinterpret_cast(_out_feats.data_ptr()); + + if (num_out_feats <= 32) + { + constexpr int G = 128; + constexpr int CTA_M = 16; + constexpr int CTA_N = 128; + constexpr int CTA_K = 128; + constexpr int WARP_M = 16; + constexpr int WARP_N = 32; + constexpr int WARP_K = 64; + constexpr int SPLITK = 2; + constexpr int STAGES = 4; + KERNEL_LAUNCH_CODE + } + else if (num_out_feats <= 64) + { + constexpr int G = 128; + constexpr int CTA_M = 16; + constexpr int CTA_N = 128; + constexpr int CTA_K = 128; + constexpr int WARP_M = 16; + constexpr int WARP_N = 32; + constexpr int WARP_K = 64; + constexpr int SPLITK = 1; + constexpr int STAGES = 3; + KERNEL_LAUNCH_CODE + } + else if (num_out_feats <= 128) + { + constexpr int G = 128; + constexpr int CTA_M = 32; + constexpr int CTA_N = 128; + constexpr int CTA_K = 128; + constexpr int WARP_M = 32; + constexpr int WARP_N = 32; + constexpr int WARP_K = 64; + constexpr int SPLITK = 1; + constexpr int STAGES = 4; + KERNEL_LAUNCH_CODE + } + else if (num_out_feats <= 192) + { + constexpr int G = 128; + constexpr int CTA_M = 64; + constexpr int CTA_N = 128; + constexpr int CTA_K = 64; + constexpr int WARP_M = 64; + constexpr int WARP_N = 32; + constexpr int WARP_K = 64; + constexpr int SPLITK = 1; + constexpr int STAGES = 4; + KERNEL_LAUNCH_CODE + } + else + { + constexpr int G = 128; + constexpr int CTA_M = 64; + constexpr int CTA_N = 128; + constexpr int CTA_K = 64; + constexpr int WARP_M = 64; + constexpr int WARP_N = 32; + constexpr int WARP_K = 64; + constexpr int STAGES = 4; + + constexpr int NUM_WARPS = (CTA_M / WARP_M) * (CTA_N / WARP_N); + constexpr int kSmemByteSize = (CTA_M * (CTA_K + SMEM_PAD_A) + CTA_N * (CTA_K + SMEM_PAD_B) / kInterleave + CTA_N) * STAGES * sizeof(ctype); + if (kSmemByteSize >= 99 * 1024) + { + printf("This kernel requires %d Bytes of shared memory, which exceeds device limit.\n", kSmemByteSize); + return _out_feats; + } + int j_factors1 = num_out_channels / CTA_N / 1; + dim3 num_blocks((num_out_feats + CTA_M - 1) / CTA_M * j_factors1); + dim3 threads_per_block(WARP_SIZE, NUM_WARPS); + auto kernel_func = gemm_w4a16_T2; + cudaFuncSetAttribute(kernel_func, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemByteSize); + kernel_func<<>>( + in_feats, kernel, scales, zeros, out_feats, num_in_feats, num_out_channels, num_in_channels); + } + }); + + return _out_feats; +} \ No newline at end of file diff --git a/llm-awq/awq/kernels/csrc/quantization_new/gemm/gemm_cuda.h b/llm-awq/awq/kernels/csrc/quantization_new/gemm/gemm_cuda.h new file mode 100644 index 0000000000000000000000000000000000000000..d5a7b158587a51e768163c5868e95f02dd8413d9 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/quantization_new/gemm/gemm_cuda.h @@ -0,0 +1,3 @@ +#include + +torch::Tensor gemm_forward_cuda_new(torch::Tensor _in_feats, torch::Tensor _kernel, torch::Tensor _scales, torch::Tensor _zeros); diff --git a/llm-awq/awq/kernels/csrc/quantization_new/gemm/semaphore.h b/llm-awq/awq/kernels/csrc/quantization_new/gemm/semaphore.h new file mode 100644 index 0000000000000000000000000000000000000000..acc636f745c53fec08521c7ec863b5d1baf675f4 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/quantization_new/gemm/semaphore.h @@ -0,0 +1,109 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Implementation of a CTA-wide semaphore for inter-CTA synchronization. +*/ + +#pragma once + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// namespace cutlass { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// CTA-wide semaphore for inter-CTA synchronization. +class Semaphore +{ +public: + int *lock; + bool wait_thread; + int state; + +public: + /// Implements a semaphore to wait for a flag to reach a given value + __host__ __device__ Semaphore(int *lock_, int thread_id) : lock(lock_), + wait_thread(thread_id < 0 || thread_id == 0), + state(-1) + { + } + + /// Permit fetching the synchronization mechanism early + __device__ void fetch() + { + if (wait_thread) + { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700 + asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n" : "=r"(state) : "l"(lock)); +#else + asm volatile("ld.global.cg.b32 %0, [%1];\n" : "=r"(state) : "l"(lock)); +#endif + } + } + + /// Gets the internal state + __device__ int get_state() const + { + return state; + } + + /// Waits until the semaphore is equal to the given value + __device__ void wait(int status = 0) + { + while (__syncthreads_and(state != status)) + { + fetch(); + } + + __syncthreads(); + } + + /// Updates the lock with the given result + __device__ void release(int status = 0) + { + __syncthreads(); + + if (wait_thread) + { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700 + asm volatile("st.global.release.gpu.b32 [%0], %1;\n" : : "l"(lock), "r"(status)); +#else + asm volatile("st.global.cg.b32 [%0], %1;\n" : : "l"(lock), "r"(status)); +#endif + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// } // namespace cutlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/llm-awq/awq/kernels/csrc/quantization_new/gemv/gemv_cuda.cu b/llm-awq/awq/kernels/csrc/quantization_new/gemv/gemv_cuda.cu new file mode 100644 index 0000000000000000000000000000000000000000..235293e43990a87bb4b1e6abb38cbc148c077478 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/quantization_new/gemv/gemv_cuda.cu @@ -0,0 +1,339 @@ +/* + * Modified from NVIDIA [TRT-LLM](https://github.com/NVIDIA/TensorRT-LLM/tree/d37b507f41a87457fe9f10f7459d08f5db235745/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv) + * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/* +@article{lin2023awq, + title={AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration}, + author={Lin, Ji and Tang, Jiaming and Tang, Haotian and Yang, Shang and Dang, Xingyu and Han, Song}, + journal={arXiv}, + year={2023} +} +*/ + +#include +#include +#include +#include "gemv_cuda.h" +#include "../dequantize.cuh" +#include "../dispatch_utils.cuh" +#define PACK_FACTOR 8 +#define WARP_SIZE 32 +#define MEM_ACCESS_SIZE 128 + +// Reduce sum within the warp using the tree reduction algorithm. +template +__device__ __forceinline__ static void warp_reduce(T* psum, float (*out_smem)[Num * 4]) +{ + // kInterleave = 4 + float fpsum[Num]; + #pragma unroll + for (int i = 0; i < Num; ++i) + { + fpsum[i] = static_cast(psum[i]); + } + + #pragma unroll + for (int i = 0; i < Num; ++i) + { + // T0 + T1 + T8 + T9 + T16 + T17 + T24 + T25 (kInterleave = 4) + fpsum[i] += __shfl_xor_sync(~0, fpsum[i], 16); + fpsum[i] += __shfl_xor_sync(~0, fpsum[i], 8); + fpsum[i] += __shfl_xor_sync(~0, fpsum[i], 1); + } + __syncthreads(); + int warp = threadIdx.x / WarpSize, lane = threadIdx.x % WarpSize; + if (lane == 0 || lane == 2 || lane == 4 || lane == 6) + { + #pragma unroll + for (int i = 0; i < Num; ++i) + { + out_smem[warp][i * 4 + lane / 2] = fpsum[i]; + } + } + __syncthreads(); +}; + +__device__ __forceinline__ int make_divisible(int c, int divisor){ + return (c + divisor - 1) / divisor; +} + +template +__global__ void gemv_kernel( + const T* inputs, const uint32_t* weight, const T* scales, const T* zeros, T* outputs, + const int IC, const int OC) +{ + const int kStride = 64; + const int kElemsPerThread = MEM_ACCESS_SIZE / 4; + const int kThreadsNumPerTile = kStride / kElemsPerThread; + // assert(MEM_ACCESS_SIZE == 128); + + using T2 = typename std::conditional< + std::is_same::value, + half2, + nv_bfloat162 + >::type; + + static constexpr int kShuffleSize = 32; + static constexpr int kShuffleBasicTile = 2; + static constexpr int kShuffleContinous = 4; + static constexpr int kShuffleStrided = 4; + + constexpr int Num = NPerBlock * Batch; + constexpr int kInterleave = 4; + + T local_inputs[kElemsPerThread]; + uint32_t local_qweights[MEM_ACCESS_SIZE / 32]; + T half_weight_buffer[kElemsPerThread]; + T dequantized_weight[kElemsPerThread * NPerBlock]; + T local_scale[NPerBlock]; + T local_scaled_zeros[NPerBlock]; + + T psum[Num]; + for (int i = 0; i < Num; ++i) + psum[i] = static_cast(0.f); + + // extern __shared__ uint8_t shmem[]; + // float(*out_smem)[Num * kInterleave] = reinterpret_cast(shmem); + __shared__ float out_smem[BlockSize / WARP_SIZE * 2][Num * kInterleave]; + + const int blk_row_offset = blockIdx.x * NPerBlock * kInterleave; + const int thd_row_offset = (threadIdx.x / kThreadsNumPerTile) % kInterleave; + const int act_k_offset = threadIdx.x / (kThreadsNumPerTile * kInterleave) * kStride + + (threadIdx.x % kThreadsNumPerTile) * kElemsPerThread; + const int group_offset = act_k_offset / GroupSize; + // TODO: use make_divisible + const uint32_t* blk_weight_ptr = weight + blk_row_offset * IC / PACK_FACTOR; + const T* scale_ptr = scales + blk_row_offset + thd_row_offset + group_offset * OC; + const T* zeros_ptr = zeros + blk_row_offset + thd_row_offset + group_offset * OC; + const T* inputs_ptr = inputs + act_k_offset; + + const int act_forward_step = BlockSize * kElemsPerThread / kInterleave; + const int scale_forward_step = act_forward_step / GroupSize * OC; + + // Main loop iteration, each block completes the outputs for several OCs + for (int kk = threadIdx.x * kElemsPerThread; kk < IC * kInterleave; kk += BlockSize * kElemsPerThread) + { + // Load qweight, scales and scaled_zeros + #pragma unroll + for (int idx = 0; idx < NPerBlock; ++idx) + { + // use float4 to load weights, each thread load 32 int4 numbers (1 x float4, 128 bit) + *((float4*)(local_qweights)) = + *((float4*)(blk_weight_ptr + (idx * kInterleave * IC + kk)/ PACK_FACTOR)); + local_scale[idx] = *(scale_ptr + idx * kInterleave); + local_scaled_zeros[idx] = *(zeros_ptr + idx * kInterleave); + + // Map int4 qweight to fp format + #pragma unroll + for (int i = 0; i < MEM_ACCESS_SIZE / 32; ++i) + { + // Converts 32 bits (8 x int4) to 8 fp16 + dequantize_s4_to_fp16x2(*reinterpret_cast(local_qweights + i), reinterpret_cast(half_weight_buffer + i * PACK_FACTOR)); + } + + // Dequantize (apply s/z) and shuffle elements to match the weight packing format + #pragma unroll + for (int i = 0; i < kShuffleContinous; ++i) + { + #pragma unroll + for (int j = 0; j < kShuffleStrided; ++j) + { + T2 w = + *reinterpret_cast( + half_weight_buffer + (i + j * kShuffleContinous)* kShuffleBasicTile + ); + if constexpr (std::is_same::value) + { + w = __hfma2(w, __half2half2(local_scale[idx]), __half2half2(local_scaled_zeros[idx])); + } + else + { + w = __hfma2(w, __bfloat162bfloat162(local_scale[idx]), __bfloat162bfloat162(local_scaled_zeros[idx])); + } + dequantized_weight[((i * kShuffleStrided + j) * kShuffleBasicTile + 0) + * NPerBlock + idx] + = w.x; + dequantized_weight[((i * kShuffleStrided + j) * kShuffleBasicTile + 1) + * NPerBlock + idx] + = w.y; + } + } + } + #pragma unroll + for (int batch_idx = 0; batch_idx < Batch; ++batch_idx) + { + const T* local_inputs_ptr = inputs_ptr + batch_idx * IC; + #pragma unroll + for (int idx = 0; idx < kElemsPerThread / 8; ++idx) + { + // load activation, 8 halves (128 bits) / step. + *((float4*)(local_inputs + idx * 8)) = *((float4*)(local_inputs_ptr + idx * 8)); + } + // Perform the MACs + #pragma unroll + for (int x = 0; x < NPerBlock / 2; ++x) + { + #pragma unroll + for (int y = 0; y < kElemsPerThread; ++y) + { + if constexpr (std::is_same::value) + { + *reinterpret_cast(psum + batch_idx * NPerBlock + x * 2) + = __hfma2(*reinterpret_cast(dequantized_weight + y * NPerBlock + x * 2), + __half2half2(local_inputs[y]), + *reinterpret_cast(psum + batch_idx * NPerBlock + x * 2)); + } + else + { + *reinterpret_cast(psum + batch_idx * NPerBlock + x * 2) + = __hfma2(*reinterpret_cast(dequantized_weight + y * NPerBlock + x * 2), + __bfloat162bfloat162(local_inputs[y]), + *reinterpret_cast(psum + batch_idx * NPerBlock + x * 2)); + } + } + } + } + inputs_ptr += act_forward_step; + scale_ptr += scale_forward_step; + zeros_ptr += scale_forward_step; + } + + warp_reduce(psum, out_smem); + + // Num * Interleave = batch * NPerBlock * Interleave -> 1 thread_block write back num + for (int i = threadIdx.x; i < Num * kInterleave; i += BlockSize) + { + int batch_idx = i / (NPerBlock * kInterleave); + int oc_idx = i % (NPerBlock * kInterleave); + float acc = 0.f; + for (int j = 0; j < BlockSize / WARP_SIZE; ++j) + { + acc += out_smem[j][i]; + } + outputs[batch_idx * OC + blk_row_offset + oc_idx] = static_cast(acc); + } +} + +/* +Computes GEMV (PyTorch interface). + +Args: + _in_feats: tensor of shape [B, IC]; + _kernel: int tensor of shape [OC, IC // 8]; + _zeros: int tensor of shape [OC, IC // G // 8]; + _scaling_factors: tensor of shape [OC, IC // G]; + blockDim_x: size of thread block, dimension x, where blockDim_x * workload_per_thread = IC; + blockDim_y: size of thread block, dimension y, where blockDim_y * gridDim_y = OC; + +Returns: + out_feats: tensor of shape [B, OC]; +*/ +torch::Tensor gemv_forward_cuda_new( + torch::Tensor _in_feats, + torch::Tensor _kernel, + torch::Tensor _scaling_factors, + torch::Tensor _zeros, + int m, + int n, + int k, + int group_size) +{ + + std::vector output_shape = _in_feats.sizes().vec(); + output_shape.back() = n; + + auto data_type = _in_feats.scalar_type(); + TORCH_CHECK(_scaling_factors.scalar_type() == data_type); + TORCH_CHECK(_zeros.scalar_type() == data_type); + + auto options = torch::TensorOptions().dtype(_in_feats.dtype()).device(_in_feats.device()); + at::Tensor _out_feats = torch::empty(output_shape, options); + + DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP16(data_type, ctype, { + auto in_feats = reinterpret_cast(_in_feats.data_ptr()); + auto kernel = reinterpret_cast(_kernel.data_ptr()); + auto zeros = reinterpret_cast(_zeros.data_ptr()); + auto scaling_factors = reinterpret_cast(_scaling_factors.data_ptr()); + auto out_feats = reinterpret_cast(_out_feats.data_ptr()); + + static constexpr int N_PER_BLOCK = 2; + static constexpr int K_INTERLEAVE = 4; + static constexpr int BLOCK_SIZE = 256; + + dim3 num_blocks(n / N_PER_BLOCK / K_INTERLEAVE); + dim3 num_threads(BLOCK_SIZE); + + // if (group_size == 64) + // { + // gemv_kernel_g64<<>>( + // // pointers + // in_feats, kernel, zeros, scaling_factors, out_feats, + // // constants + // num_in_channels, num_out_channels + // ); + // } + if (group_size == 128) + { + switch (m) + { + case 1: + gemv_kernel<<>>( + in_feats, kernel, scaling_factors, zeros, out_feats, k, n + ); + break; + case 2: + gemv_kernel<<>>( + in_feats, kernel, scaling_factors, zeros, out_feats, k, n + ); + break; + case 3: + gemv_kernel<<>>( + in_feats, kernel, scaling_factors, zeros, out_feats, k, n + ); + break; + case 4: + gemv_kernel<<>>( + in_feats, kernel, scaling_factors, zeros, out_feats, k, n + ); + break; + case 5: + gemv_kernel<<>>( + in_feats, kernel, scaling_factors, zeros, out_feats, k, n + ); + break; + case 6: + gemv_kernel<<>>( + in_feats, kernel, scaling_factors, zeros, out_feats, k, n + ); + break; + case 7: + gemv_kernel<<>>( + in_feats, kernel, scaling_factors, zeros, out_feats, k, n + ); + break; + default: + throw std::runtime_error("Unsupported batch size for gemv kernel.\n"); + } + } + else + { + throw std::runtime_error("Unsupported group size for gemv kernel.\n"); + } + }); + return _out_feats; +} + diff --git a/llm-awq/awq/kernels/csrc/quantization_new/gemv/gemv_cuda.h b/llm-awq/awq/kernels/csrc/quantization_new/gemv/gemv_cuda.h new file mode 100644 index 0000000000000000000000000000000000000000..181637cfc8716a9293ebea721593d4f121d2dfe7 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/quantization_new/gemv/gemv_cuda.h @@ -0,0 +1,12 @@ +#pragma once +#include + +torch::Tensor gemv_forward_cuda_new( + torch::Tensor _in_feats, + torch::Tensor _kernel, + torch::Tensor _scaling_factors, + torch::Tensor _zeros, + int m, + int n, + int k, + int group_size); diff --git a/llm-awq/awq/kernels/csrc/rope_new/fused_rope_with_pos.cu b/llm-awq/awq/kernels/csrc/rope_new/fused_rope_with_pos.cu new file mode 100644 index 0000000000000000000000000000000000000000..b81ebff6da00a691de7e001d03562a9477d6f8a6 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/rope_new/fused_rope_with_pos.cu @@ -0,0 +1,407 @@ +// Modified from https://github.com/NVIDIA/TransformerEngine +// Modified by Shang Yang. + +/************************************************************************* + * Copyright (c) 2022-2024, NVIDIA CORPORATION & AFFILIATES. All rights + *reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include +#include +#include +#include +#include + +#include "fused_rope_with_pos.h" +// #include + +// #include "../common.h" +// #include "../util/logging.h" +// #include "../utils.cuh" +#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \ + AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) + +#define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \ + AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__)) + +#define THREADS_PER_WARP 32 + +template +__device__ void fused_rope_with_pos_block_forward( + const scalar_t *src, const float *freqs, scalar_t *dst, + const int offset_block, const int offset_block_dst, const int h, + const int d, const int d2, const int stride_h, const int stride_d, + const int o_stride_h, const int o_stride_d) { + int s_id = blockIdx.x; + int s = gridDim.x; + int b_id = blockIdx.y; +#pragma unroll + for (int d_id = threadIdx.x; d_id < d2; d_id += blockDim.x) { + float v_cos, v_sin; + sincosf(freqs[(b_id * s + s_id) * d2 + d_id], &v_sin, &v_cos); +#pragma unroll + for (int h_id = threadIdx.y; h_id < h; h_id += blockDim.y) { + int offset_src = offset_block + h_id * stride_h + d_id * stride_d; + int offset_dst = offset_block_dst + h_id * o_stride_h + d_id * o_stride_d; + float v_src = src[offset_src]; + float v_src_rotate = + (d_id + d2 / 2 < d2) + ? -static_cast(src[offset_src + (d2 / 2) * stride_d]) + : static_cast(src[offset_src + (d2 / 2 - d2) * stride_d]); + dst[offset_dst] = v_src * v_cos + v_src_rotate * v_sin; + } + } + + // copy the rest + if (d > d2) { +#pragma unroll + for (int h_id = threadIdx.y; h_id < h; h_id += blockDim.y) { + int offset_head = offset_block + h_id * stride_h; + int offset_head_dst = offset_block_dst + h_id * o_stride_h; +#pragma unroll + for (int d_id = d2 + threadIdx.x; d_id < d; d_id += blockDim.x) { + dst[offset_head_dst + d_id * o_stride_d] = + src[offset_head + d_id * stride_d]; + } + } + } +} + +// template +// __device__ void fused_rope_block_backward(const scalar_t *src, const float +// *freqs, scalar_t *dst, +// const int offset_block, const int +// offset_block_dst, const int h, +// const int d, const int d2, const +// int stride_h, const int stride_d, +// const int o_stride_h, const int +// o_stride_d) { +// int s_id = blockIdx.x; +// #pragma unroll +// for (int d_id = threadIdx.x; d_id < d2; d_id += blockDim.x) { +// float v_cos = cosf(freqs[s_id * d2 + d_id]); +// float v_sin = (d_id + d2 / 2 < d2) ? sinf(freqs[s_id * d2 + d_id + d2 / +// 2]) +// : -sinf(freqs[s_id * d2 + d_id + d2 / +// 2 - d2]); +// #pragma unroll +// for (int h_id = threadIdx.y; h_id < h; h_id += blockDim.y) { +// int offset_src = offset_block + h_id * stride_h + d_id * stride_d; +// int offset_dst = offset_block_dst + h_id * o_stride_h + d_id * +// o_stride_d; float v_src = src[offset_src]; float v_src_rotate = (d_id + +// d2 / 2 < d2) ? src[offset_src + (d2 / 2) * stride_d] +// : src[offset_src + (d2 / 2 - +// d2) * stride_d]; +// dst[offset_dst] = v_src * v_cos + v_src_rotate * v_sin; +// } +// } + +// // handle the tail +// if (d > d2) { +// #pragma unroll +// for (int h_id = threadIdx.y; h_id < h; h_id += blockDim.y) { +// int offset_head = offset_block + h_id * stride_h; +// int offset_head_dst = offset_block_dst + h_id * o_stride_h; +// #pragma unroll +// for (int d_id = d2 + threadIdx.x; d_id < d; d_id += blockDim.x) { +// dst[offset_head_dst + d_id * o_stride_d] = src[offset_head + d_id * +// stride_d]; +// } +// } +// } +// } + +template +__global__ void fused_rope_with_pos_forward_kernel( + const scalar_t *src, const float *freqs, scalar_t *dst, const int h, + const int d, const int d2, const int stride_s, const int stride_b, + const int stride_h, const int stride_d, const int o_stride_s, + const int o_stride_b, const int o_stride_h, const int o_stride_d) { + int s_id = blockIdx.x, b_id = blockIdx.y; + int offset_block = s_id * stride_s + b_id * stride_b; + int offset_block_dst = s_id * o_stride_s + b_id * o_stride_b; + fused_rope_with_pos_block_forward( + src, freqs, dst, offset_block, offset_block_dst, h, d, d2, stride_h, + stride_d, o_stride_h, o_stride_d); +} + +// template +// __global__ void fused_rope_backward_kernel(const scalar_t *src, const float +// *freqs, scalar_t *dst, +// const int h, const int d, const +// int d2, const int stride_s, const +// int stride_b, const int stride_h, +// const int stride_d, const int +// o_stride_s, const int o_stride_b, +// const int o_stride_h, const int +// o_stride_d) { +// int s_id = blockIdx.x, b_id = blockIdx.y; +// int offset_block = s_id * stride_s + b_id * stride_b; +// int offset_block_dst = s_id * o_stride_s + b_id * o_stride_b; +// fused_rope_block_backward(src, freqs, dst, offset_block, +// offset_block_dst, h, d, d2, stride_h, +// stride_d, o_stride_h, o_stride_d); +// } + +template +void fused_rope_with_pos_forward_launcher( + const scalar_t *input, const float *freqs, scalar_t *output, const int s, + const int b, const int h, const int d, const int d2, const int stride_s, + const int stride_b, const int stride_h, const int stride_d, + const int o_stride_s, const int o_stride_b, const int o_stride_h, + const int o_stride_d, cudaStream_t stream) { + int warps_per_block = h < 16 ? 4 : 8; + dim3 blocks(s, b); + dim3 threads(THREADS_PER_WARP, warps_per_block); + + fused_rope_with_pos_forward_kernel<<>>( + input, freqs, output, h, d, d2, stride_s, stride_b, stride_h, stride_d, + o_stride_s, o_stride_b, o_stride_h, o_stride_d); + // NVTE_CHECK_CUDA(cudaGetLastError()); +} + +// template +// void fused_rope_backward_launcher(const scalar_t *output_grads, const float +// *freqs, +// scalar_t *input_grads, const int s, const +// int b, const int h, const int d, const int +// d2, const int stride_s, const int stride_b, +// const int stride_h, const int stride_d, +// const int o_stride_s, const int o_stride_b, +// const int o_stride_h, const int o_stride_d, +// cudaStream_t stream) { +// int warps_per_block = h < 16 ? 4 : 8; +// dim3 blocks(s, b); +// dim3 threads(THREADS_PER_WARP, warps_per_block); + +// fused_rope_backward_kernel<<>>( +// output_grads, freqs, input_grads, h, d, d2, stride_s, stride_b, +// stride_h, stride_d, o_stride_s, o_stride_b, o_stride_h, o_stride_d); +// // NVTE_CHECK_CUDA(cudaGetLastError()); +// } + +template +void fused_rope_with_pos_forward(const at::Tensor &input, + const at::Tensor &freqs, at::Tensor &output, + const int s, const int b, const int h, + const int d, const int d2, const int stride_s, + const int stride_b, const int stride_h, + const int stride_d, const int o_stride_s, + const int o_stride_b, const int o_stride_h, + const int o_stride_d, cudaStream_t stream) { + // TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT( + // input.data.dtype, scalar_t, + fused_rope_with_pos_forward_launcher( + reinterpret_cast(input.data_ptr()), + reinterpret_cast(freqs.data_ptr()), + reinterpret_cast(output.data_ptr()), s, b, h, d, d2, stride_s, + stride_b, stride_h, stride_d, o_stride_s, o_stride_b, o_stride_h, + o_stride_d, stream); + // ); +} + +// template +// void fused_rope_backward(const at::Tensor &output_grads, const at::Tensor +// &freqs, at::Tensor &input_grads, +// const int s, const int b, const int h, const int d, +// const int d2, const int stride_s, const int +// stride_b, const int stride_h, const int stride_d, +// const int o_stride_s, const int o_stride_b, const +// int o_stride_h, const int o_stride_d, cudaStream_t +// stream) { +// // TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT( +// // output_grads.data.dtype, scalar_t, +// fused_rope_backward_launcher(reinterpret_cast(output_grads.data_ptr()), +// reinterpret_cast(freqs.data_ptr()), +// reinterpret_cast(input_grads.data_ptr()), s, b, h, d, +// d2, stride_s, stride_b, stride_h, +// stride_d, o_stride_s, o_stride_b, +// o_stride_h, o_stride_d, stream); +// // ); +// } + +template +void nvte_fused_rope_with_pos_forward( + const at::Tensor input, const at::Tensor freqs, at::Tensor output, + const int s, const int b, const int h, const int d, const int d2, + const int stride_s, const int stride_b, const int stride_h, + const int stride_d, const int o_stride_s, const int o_stride_b, + const int o_stride_h, const int o_stride_d, cudaStream_t stream) { + // NVTE_API_CALL(nvte_fused_rope_forward); + // using namespace transformer_engine; + fused_rope_with_pos_forward( + input, freqs, output, s, b, h, d, d2, stride_s, stride_b, stride_h, + stride_d, o_stride_s, o_stride_b, o_stride_h, o_stride_d, stream); +} + +// template +// void nvte_fused_rope_backward(const at::Tensor output_grads, const at::Tensor +// freqs, +// at::Tensor input_grads, const int s, const int +// b, const int h, const int d, const int d2, +// const int stride_s, const int stride_b, const +// int stride_h, const int stride_d, const int +// o_stride_s, const int o_stride_b, const int +// o_stride_h, const int o_stride_d, cudaStream_t +// stream) { +// // NVTE_API_CALL(nvte_fused_rope_backward); +// // using namespace transformer_engine; +// fused_rope_backward(output_grads, freqs, input_grads, s, b, h, d, +// d2, stride_s, stride_b, +// stride_h, stride_d, o_stride_s, o_stride_b, o_stride_h, +// o_stride_d, stream); +// } + +// Interface for Python +at::Tensor fused_rope_with_pos_forward_func( + const at::Tensor &input, const at::Tensor &freqs, + const bool transpose_output_memory) { + // using namespace transformer_engine; + // TORCH_CHECK(input.dim() == 4, "expected 4D tensor"); + // TORCH_CHECK(freqs.dim() == 4, "expected 4D tensor"); + // TORCH_CHECK(input.size(0) <= freqs.size(0), + // "expected freqs tensor has a longer sequence length than + // input"); + // TORCH_CHECK(freqs.size(1) == 1 && freqs.size(2) == 1, + // "expected the second and third dims of the freqs tensor equal + // 1"); + // TORCH_CHECK(input.size(3) >= freqs.size(3), + // "expected the last dim of the input tensor equals or is " + // "greater than the freqs tensor"); + // TORCH_CHECK(freqs.scalar_type() == at::ScalarType::Float, + // "Dtype of the freqs tensor must be float"); + + // input sizes: (s, b, h, d) + // s: sequence length + // b: batch size + // h: head num + // d: dim of each head + const int s = input.size(0); + const int b = input.size(1); + const int h = input.size(2); + const int d = input.size(3); + // input strides + const int stride_s = input.stride(0); + const int stride_b = input.stride(1); + const int stride_h = input.stride(2); + const int stride_d = input.stride(3); + // freqs' shape is always (s, 1, 1, d2), so the strides are same under + // different memory formats + // freqs' shape is now (B, S, D) + const int d2 = freqs.size(-1); + + // output + auto act_options = input.options().requires_grad(false); + at::Tensor output; + if (transpose_output_memory) { + output = torch::empty({b, s, h, d}, act_options).transpose(0, 1); + } else { + output = torch::empty({s, b, h, d}, act_options); + } + // output strides + const int o_stride_s = output.stride(0); + const int o_stride_b = output.stride(1); + const int o_stride_h = output.stride(2); + const int o_stride_d = output.stride(3); + + auto input_cu = input; + auto freqs_cu = freqs; + auto output_cu = output; + + VLLM_DISPATCH_FLOATING_TYPES( + input.scalar_type(), "nvte_fused_rope_forward", [&] { + nvte_fused_rope_with_pos_forward( + input_cu.data(), freqs_cu.data(), output_cu.data(), s, b, h, d, d2, + stride_s, stride_b, stride_h, stride_d, o_stride_s, o_stride_b, + o_stride_h, o_stride_d, at::cuda::getCurrentCUDAStream()); + }); + + // nvte_fused_rope_forward(input_cu.data(), + // freqs_cu.data(), output_cu.data(), s, b, h, d, d2, + // stride_s, stride_b, stride_h, stride_d, o_stride_s, + // o_stride_b, o_stride_h, o_stride_d, + // at::cuda::getCurrentCUDAStream()); + + return output; +} + +// // Interface for Python +// at::Tensor fused_rope_backward_func(const at::Tensor &output_grads, const +// at::Tensor &freqs, +// const bool transpose_output_memory) { +// // using namespace transformer_engine; +// // TORCH_CHECK(output_grads.dim() == 4, "expected 4D tensor"); +// // TORCH_CHECK(freqs.dim() == 4, "expected 4D tensor"); +// // TORCH_CHECK(output_grads.size(0) <= freqs.size(0), +// // "expected freqs tensor has a longer sequence length than +// output_grads"); +// // TORCH_CHECK(freqs.size(1) == 1 && freqs.size(2) == 1, +// // "expected the second and third dims of the freqs tensor +// equal 1"); +// // TORCH_CHECK(output_grads.size(3) >= freqs.size(3), +// // "expected the last dim of the output_grads tensor equals or +// is " +// // "greater than the freqs tensor"); +// // TORCH_CHECK(freqs.scalar_type() == at::ScalarType::Float, +// // "Dtype of the freqs tensor must be float"); + +// // output_grads sizes: (s, b, h, d) +// // s: sequence length +// // b: batch size +// // h: head num +// // d: dim of each head +// const int s = output_grads.size(0); +// const int b = output_grads.size(1); +// const int h = output_grads.size(2); +// const int d = output_grads.size(3); +// // output_grads strides +// const int stride_s = output_grads.stride(0); +// const int stride_b = output_grads.stride(1); +// const int stride_h = output_grads.stride(2); +// const int stride_d = output_grads.stride(3); +// // freqs' shape is always (s, 1, 1, d2), so the strides are same under +// // different memory formats +// const int d2 = freqs.size(3); + +// auto act_options = output_grads.options().requires_grad(false); +// at::Tensor input_grads; +// if (transpose_output_memory) { +// input_grads = torch::empty({b, s, h, d}, act_options).transpose(0, 1); +// } else { +// input_grads = torch::empty({s, b, h, d}, act_options); +// } +// const int o_stride_s = input_grads.stride(0); +// const int o_stride_b = input_grads.stride(1); +// const int o_stride_h = input_grads.stride(2); +// const int o_stride_d = input_grads.stride(3); + +// auto output_grads_cu = output_grads; +// auto freqs_cu = freqs; +// auto input_grads_cu = input_grads; + +// VLLM_DISPATCH_FLOATING_TYPES( +// output_grads.scalar_type(), "nvte_fused_rope_forward", [&] { +// nvte_fused_rope_backward(output_grads_cu.data(), +// freqs_cu.data(), input_grads_cu.data(), s, b, h, +// d, d2, stride_s, stride_b, stride_h, +// stride_d, o_stride_s, o_stride_b, +// o_stride_h, o_stride_d, +// at::cuda::getCurrentCUDAStream()); +// }); + +// // nvte_fused_rope_backward(output_grads_cu.data(), freqs_cu.data(), +// input_grads_cu.data(), s, b, h, +// // d, d2, stride_s, stride_b, stride_h, stride_d, +// o_stride_s, o_stride_b, +// // o_stride_h, o_stride_d, +// at::cuda::getCurrentCUDAStream()); + +// return input_grads; +// } diff --git a/llm-awq/awq/kernels/csrc/rope_new/fused_rope_with_pos.h b/llm-awq/awq/kernels/csrc/rope_new/fused_rope_with_pos.h new file mode 100644 index 0000000000000000000000000000000000000000..6f999d186e1f2ed76f8c31878b8591b7b94d12ab --- /dev/null +++ b/llm-awq/awq/kernels/csrc/rope_new/fused_rope_with_pos.h @@ -0,0 +1,5 @@ +#include + +at::Tensor fused_rope_with_pos_forward_func(const at::Tensor &input, + const at::Tensor &freqs, + const bool transpose_output_memory); diff --git a/llm-awq/awq/kernels/csrc/w8a8/act.cu b/llm-awq/awq/kernels/csrc/w8a8/act.cu new file mode 100644 index 0000000000000000000000000000000000000000..5dce6c1c55d7ba33c94094a97a6dc09ee6aeee59 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/w8a8/act.cu @@ -0,0 +1,141 @@ +#include +#include +#include + +#include "dispatch_utils.h" +#include "utils.cuh" +#include "reduction_utils.cuh" + +namespace vllm { + +template __device__ __forceinline__ T silu(const T &x) { + // x * sigmoid(x) + return (T)(((float)x) / (1.0f + expf((float)-x))); +} + +template __device__ __forceinline__ T gelu_new(const T &x) { + const half x3 = (half)(x * x * x); + const T t = (T)tanhf((T)((T)0.79788456f * (half)(x + (T)((T)0.044715f * x3)))); + return ((T)0.5) * x * (((T)1.0) + t); +} + +template +__device__ __forceinline__ T gelu_fast(const T &x) { + const half f = (half)x; + const T t = + (T)tanhf(((T)(f * (T)0.79788456f)) * (((T)1.0) + (T)((T)0.044715f * f) * x)); + return ((T)0.5) * x * (((T)1.0) + t); +} + + + +// dequant int32 input, apply silu and mul, then per token quant to int8 +template +__global__ void gelu_and_quant_kernel( + int8_t *__restrict__ out, // [..., d] + half *__restrict__ input, // [..., d] + const int d, + scale_type * scale_out, // [num_tokens] + half *__restrict__ tmp = nullptr // [num_tokens, d] +) { + const int token_idx = blockIdx.x; + const float max_value= 127.0f; + if constexpr (use_per_token_quant) { + float amax_val = 0.0f; + const half zero = 0.0001f; + + for (int idx = threadIdx.x; idx < d; idx += blockDim.x) { + const half x = + (half)__ldg(&input[token_idx * d + idx]); + half t = gelu_fast(x); + tmp[token_idx * d + idx] = t; + t = t > zero ? t : -t; + if ((float)t > amax_val) + amax_val = (float)t; + } + + __shared__ float s_amax; + const float block_amax_val = blockReduceMax(amax_val); + if (threadIdx.x == 0) { + s_amax = block_amax_val; + scale_out[token_idx] = half(block_amax_val / max_value); + } + __syncthreads(); + + float tmp_scale = max_value / s_amax; + for (int idx = threadIdx.x; idx < d; idx += blockDim.x) { + out[token_idx * d + idx] = + float_to_int8_rn((half)tmp_scale * tmp[token_idx * d + idx]); + } + } else { + for (int idx = threadIdx.x; idx < d; idx += blockDim.x) { + const float x = + (float)__ldg(&input[token_idx * d + idx]); + out[token_idx * d + idx] = float_to_int8_rn((half)gelu_fast(x) / scale_out[0]); + } + } +} +} // namespace vllm + + + +void gelu_and_quant( + torch::Tensor &out, // [..., d] + torch::Tensor &input, // [..., d] + torch::Tensor &scale_out, // [...] + torch::Tensor &tmp // [num_tokens, d] + ) { + int64_t num_tokens = input.numel() / input.size(-1); + int d = input.size(-1); + dim3 grid(num_tokens); + dim3 block(std::min(d, 128)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + vllm::gelu_and_quant_kernel<<>>( + out.data_ptr(), reinterpret_cast(input.data_ptr()), d, reinterpret_cast(scale_out.data_ptr()),reinterpret_cast(tmp.data_ptr())); +} + + + +namespace vllm { + +template +__global__ void silu_and_mul_kernel( + scalar_t* __restrict__ out, // [..., d] + const scalar_t* __restrict__ input, // [..., 2 * d] + const int d) { + + const int token_idx = blockIdx.x; + const int64_t token_idx_d = token_idx * int64_t(d); + const int64_t token_idx_2d = token_idx_d * 2; + for (int idx = threadIdx.x; idx < d; idx += blockDim.x) { + const scalar_t x = __ldg(&input[token_idx_2d + idx]); + const scalar_t y = __ldg(&input[token_idx_2d + d + idx]); + out[token_idx_d + idx] = silu(x) * y; + } +} +} // namespace vllm + + + +torch::Tensor silu_and_mul( + torch::Tensor& input) // [..., 2 * d] +{ + int64_t num_tokens = input.numel() / input.size(-1); + int d = input.size(-1) / 2; + + std::vector output_shape = input.sizes().vec(); + output_shape[output_shape.size() - 1]=d; + auto options = + torch::TensorOptions().dtype(input.dtype()).device(input.device()); + at::Tensor output = torch::empty(output_shape, options); + + + dim3 grid(num_tokens); + dim3 block(std::min(d, 256)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "silu_and_mul_kernel", [&] { + vllm::silu_and_mul_kernel<<>>( + output.data_ptr(), input.data_ptr(), d); + }); + return output; +} \ No newline at end of file diff --git a/llm-awq/awq/kernels/csrc/w8a8/act.h b/llm-awq/awq/kernels/csrc/w8a8/act.h new file mode 100644 index 0000000000000000000000000000000000000000..98c92c2cfbaf8929321b921789cf48405d10f7a0 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/w8a8/act.h @@ -0,0 +1,29 @@ +// Inspired by TRT-LLM. +// Modified by Shang Yang and Haotian Tang. +// @article{lin2024awq, +// title={AWQ: Activation-aware Weight Quantization for On-Device LLM Compression and Acceleration}, +// author={Lin, Ji and Tang, Jiaming and Tang, Haotian and Yang, Shang and Chen, Wei-Ming and Wang, Wei-Chen and Xiao, Guangxuan and Dang, Xingyu and Gan, Chuang and Han, Song}, +// journal={Proceedings of Machine Learning and Systems}, +// volume={6}, +// pages={87--100}, +// year={2024} +// } + +#include +#include +// Inspired by vLLM-SmoothQuant: https://github.com/vllm-project/vllm/pull/1112. +#include + + +void gelu_and_quant(torch::Tensor &out, // [..., d] + torch::Tensor &input, // [..., d] + torch::Tensor &scale_out, // [num_tokens] + torch::Tensor &tmp // [num_tokens, d] +); + +torch::Tensor silu_and_mul(torch::Tensor &input // [..., 2 * d] +); + + + + diff --git a/llm-awq/awq/kernels/csrc/w8a8/dispatch_utils.h b/llm-awq/awq/kernels/csrc/w8a8/dispatch_utils.h new file mode 100644 index 0000000000000000000000000000000000000000..7c0c49d392a9806f1da29f0df61d362ed7170f59 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/w8a8/dispatch_utils.h @@ -0,0 +1,14 @@ +/* + * Adapted from + * https://github.com/pytorch/pytorch/blob/v2.0.1/aten/src/ATen/Dispatch.h + */ +#include + +#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \ + AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) + +#define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \ + AT_DISPATCH_SWITCH( \ + TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__)) diff --git a/llm-awq/awq/kernels/csrc/w8a8/layernorm.cu b/llm-awq/awq/kernels/csrc/w8a8/layernorm.cu new file mode 100644 index 0000000000000000000000000000000000000000..510d4e1e471fe67df5dd37656f23163a70fc2bf6 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/w8a8/layernorm.cu @@ -0,0 +1,232 @@ +// Inspired by QServe https://github.com/mit-han-lab/qserve/tree/main. +// Modified by Yuming Lou. +// @article{lin2024awq, +// title={AWQ: Activation-aware Weight Quantization for On-Device LLM Compression and Acceleration}, +// author={Lin, Ji and Tang, Jiaming and Tang, Haotian and Yang, Shang and Chen, Wei-Ming and Wang, Wei-Chen and Xiao, Guangxuan and Dang, Xingyu and Gan, Chuang and Han, Song}, +// journal={Proceedings of Machine Learning and Systems}, +// volume={6}, +// pages={87--100}, +// year={2024} +// } +#include +#include +#include "dispatch_utils.h" +#include "utils.cuh" +#include "reduction_utils.cuh" + + +namespace vllm { + +// from TRTLLM +template +__inline__ __device__ Tf compute_layernorm(Tf val, float s_mean, float s_variance, const T* gamma, const T* beta, int i) +{ + Tf ret = (val - s_mean) * s_variance * cuda_cast(gamma[i]); + if (beta != nullptr) + { + ret = ret + cuda_cast(beta[i]); + } + return ret; +} + +// from TRTLLM +/* Computes the layernorm https://pytorch.org/docs/stable/generated/torch.nn.LayerNorm.html + * normed_output <- ( (input - E[input]) / Sqrt(Var[input] + eps) ) * gamma + beta + * input is [tokens, hidden_dim]. Mean and Variance are per-row (i.e. per-token) + * + * One CTA handles one row. + * + * with USE_DIFF_OF_SQUARES set to false: + * First pass (loop) computes the mean. + * Second computes the variance via Var[x] = E[(x - E[x])²]. + * Third pass computes and writes normed_output + * For better speedup, we set USE_DIFF_OF_SQUARES to true (may be faster but less accurate): + * It turns out the accuracy dosen't drop. + * First pass (loop) computes the mean and variance via Var[x] = E[x²] - E[x]² + * Second pass computes and writes normed_output + * + * + * use_shmem controls if we cache input values into shared memory + * + * Optional: with dynamic scaling, the last pass doesn't write immediately but finds the + * amax per row. A final pass scales to int8 accordingly, and writes output to + * normed_output_quant. + */ +template +__global__ void generalLayerNorm(const T* input, const T* gamma, const T* beta, T* normed_output, const float eps, + int tokens, int hidden_dim, const scale_type* scale_orig_quant_per_tensor, scale_type* scale_orig_quant_per_token, + int8_t* normed_output_quant, bool use_shmem) +{ + constexpr auto num_elems_T = num_elems::value; + using int8_packed_t = typename packed_as::type; + using float_packed_t = typename packed_as::type; + using T_scalar = typename packed_as::type; + + extern __shared__ __align__(sizeof(float)) char _shmem[]; + T* shmem = reinterpret_cast(_shmem); + __shared__ float s_mean; + __shared__ float s_variance; + + const int tidx = threadIdx.x; + const int bidx = blockIdx.x; + + float mean = 0.0f; + float variance = 0.0f; + float local_sum = 0.0f; + float local_var_sum = 0.0f; + const int n_elems = hidden_dim / num_elems_T; + for (int i = tidx; i < n_elems; i += blockDim.x) + { + const T val = input[bidx * n_elems + i]; + if (use_shmem) + { + shmem[i] = val; + } + const float_packed_t val_f = cuda_cast(val); + local_sum += cuda_sum(val_f); + if (USE_DIFF_OF_SQUARES) + { + local_var_sum += cuda_sum(val_f * val_f); + } + } + //Compute mean + if (USE_DIFF_OF_SQUARES) + { + float packed[2] = {local_sum, local_var_sum}; + blockReduceSumV2(packed); + mean = packed[0]; + variance = packed[1]; + } + else + { + mean = blockReduceSum(local_sum); + } + + if (threadIdx.x == 0) + { + mean = mean / hidden_dim; + s_mean = mean; + if (USE_DIFF_OF_SQUARES) + { + variance = (variance / hidden_dim) - (mean * mean); // Var[x] = E[x²] - E[x]² + s_variance = rsqrtf(variance + eps); + } + } + __syncthreads(); + + + if (!USE_DIFF_OF_SQUARES) + { + for (int i = tidx; i < n_elems; i += blockDim.x) + { + const T val = use_shmem ? shmem[i] : input[bidx * n_elems + i]; + float_packed_t diff = cuda_cast(val); // - s_mean; + local_var_sum += cuda_sum(diff * diff); + } + variance = blockReduceSum(local_var_sum); + + if (threadIdx.x == 0) + { + s_variance = rsqrtf(variance / hidden_dim + eps); + } + __syncthreads(); + } + + // Compute LN and Quantize + const bool with_per_token_scaling = scale_orig_quant_per_token != nullptr; + const bool with_per_tensor_scaling = scale_orig_quant_per_tensor != nullptr; + const float_packed_t scale_orig_quant + = cuda_cast(with_per_tensor_scaling ? __half2float(*scale_orig_quant_per_tensor) : 0.0f); + T_scalar amax = 1e-6f; + + for (int i = tidx; i < n_elems; i += blockDim.x) + { + const int index = bidx * n_elems + i; + const float_packed_t val_f = cuda_cast(use_shmem ? shmem[i] : input[index]); + const T val = cuda_cast(compute_layernorm(val_f, s_mean, s_variance, gamma, beta, i)); + + if (with_per_token_scaling) + { + amax = cuda_max(cuda_max(cuda_abs(val)), amax); + if (use_shmem) + { + shmem[i] = val; + } + } + else if (with_per_tensor_scaling) + { + reinterpret_cast(normed_output_quant)[index] + = cuda_cast(cuda_cast(val) * scale_orig_quant); + } + else + { + normed_output[index] = val; + } + } + + if (with_per_token_scaling) + { + float abs_max_f = blockAllReduceMax(cuda_cast(amax)); + const float dynamic_per_token_scale = 127.f / abs_max_f; + for (int i = tidx; i < n_elems; i += blockDim.x) + { + const int index = bidx * n_elems + i; + float_packed_t val_f = cuda_cast(use_shmem ? shmem[i] : input[index]); + if (!use_shmem) + { + val_f = compute_layernorm(val_f, s_mean, s_variance, gamma, beta, i); + } + + reinterpret_cast(normed_output_quant)[index] + = cuda_cast(val_f * cuda_cast(dynamic_per_token_scale)); + } + if (tidx == 0) + { + scale_orig_quant_per_token[bidx] = abs_max_f / 127.f; + } + } +} + + +} // namespace vllm + +void rms_norm_general(torch::Tensor &out, // [..., hidden_size] + torch::Tensor &input, // [..., hidden_size] + torch::Tensor &weight, // [hidden_size] + torch::Tensor &bias, // [hidden_size] + torch::Tensor &scaling, // [tokens] or [1] + float epsilon, + bool use_per_token_quant = true) { + int hidden_size = input.size(-1); + int num_tokens = input.numel() / hidden_size; + dim3 grid(num_tokens); + dim3 block(std::min(hidden_size, 128));//Reduce the idle probability of threads + block.x = 32 * ((block.x + 31) / 32); + + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "generalLayerNorm", [&] { + using T = typename FloatTypeConverter::Type; + if (use_per_token_quant) { + // per-token + vllm::generalLayerNorm<<>>( + reinterpret_cast(input.data_ptr()), + reinterpret_cast(weight.data_ptr()), + reinterpret_cast(bias.data_ptr()), + nullptr, epsilon, num_tokens, hidden_size, nullptr, scaling.data_ptr(), + out.data_ptr(), false + ); + // input, gamma, beta, normed_output, eps, tokens, hidden_dim, per_tensor_scale, per_token_scale + // normed_output_quant, use_shmem + // out.data_ptr(), input.data_ptr(), + // weight.data_ptr(), epsilon, num_tokens, hidden_size); + } else { + // per-tensor + vllm::generalLayerNorm<<>>( + reinterpret_cast(input.data_ptr()), + reinterpret_cast(weight.data_ptr()), nullptr, + nullptr, epsilon, num_tokens, hidden_size, scaling.data_ptr(), nullptr, + out.data_ptr(), false + ); + } + }); +} \ No newline at end of file diff --git a/llm-awq/awq/kernels/csrc/w8a8/layernorm.h b/llm-awq/awq/kernels/csrc/w8a8/layernorm.h new file mode 100644 index 0000000000000000000000000000000000000000..9e5d740520cf3b1f1b910e61e59ffe0dc570d810 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/w8a8/layernorm.h @@ -0,0 +1,21 @@ +// Inspired by TRT-LLM. +// Modified by Shang Yang and Haotian Tang. +// @article{lin2024awq, +// title={AWQ: Activation-aware Weight Quantization for On-Device LLM Compression and Acceleration}, +// author={Lin, Ji and Tang, Jiaming and Tang, Haotian and Yang, Shang and Chen, Wei-Ming and Wang, Wei-Chen and Xiao, Guangxuan and Dang, Xingyu and Gan, Chuang and Han, Song}, +// journal={Proceedings of Machine Learning and Systems}, +// volume={6}, +// pages={87--100}, +// year={2024} +// } + +#include +#include +void rms_norm_general(torch::Tensor &out, // [..., hidden_size] + torch::Tensor &input, // [..., hidden_size] + torch::Tensor &weight, // [hidden_size] + torch::Tensor &bias, // [hidden_size] + torch::Tensor &scaling, // [tokens] or [1] + float epsilon, + bool use_per_token_quant); + diff --git a/llm-awq/awq/kernels/csrc/w8a8/quantization.cu b/llm-awq/awq/kernels/csrc/w8a8/quantization.cu new file mode 100644 index 0000000000000000000000000000000000000000..43d02507448645e6441c7b131b69399dcad6ad3b --- /dev/null +++ b/llm-awq/awq/kernels/csrc/w8a8/quantization.cu @@ -0,0 +1,113 @@ +// Inspired by vLLM-SmoothQuant: https://github.com/vllm-project/vllm/pull/1112 and TensorRT-LLM. +// Modified by Shang Yang and Haotian Tang. +// @article{lin2024awq, +// title={AWQ: Activation-aware Weight Quantization for On-Device LLM Compression and Acceleration}, +// author={Lin, Ji and Tang, Jiaming and Tang, Haotian and Yang, Shang and Chen, Wei-Ming and Wang, Wei-Chen and Xiao, Guangxuan and Dang, Xingyu and Gan, Chuang and Han, Song}, +// journal={Proceedings of Machine Learning and Systems}, +// volume={6}, +// pages={87--100}, +// year={2024} +// } +#include +#include + +#include "utils.cuh" +#include +#include +#include "quantization.h" + +#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \ + AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) +#define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__)) + +template +__inline__ __device__ T warpReduceMax(T val) +{ +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) + val = max(val, __shfl_xor_sync(0xffffffff, val, mask, 32)); + return val; +} + +/* Calculate the maximum of all elements in a block */ +template +__inline__ __device__ T blockReduceMax(T val) +{ + static __shared__ T shared[32]; + int lane = threadIdx.x & 0x1f; // in-warp idx + int wid = threadIdx.x >> 5; // warp idx + val = warpReduceMax(val); // get maxx in each warp + if (lane == 0) // record in-warp maxx by warp Idx + shared[wid] = val; + __syncthreads(); + // Modify from blockDim.x << 5 to blockDim.x / 32. to prevent + // blockDim.x is not divided by 32 + val = (threadIdx.x < (blockDim.x / 32.f)) ? shared[lane] : -1e20f; + val = warpReduceMax(val); + return val; +} + + + +namespace vllm { +template +__global__ void quant_kernel(const T *__restrict__ input, + int8_t *__restrict__ output, scale_type scale, + int num_tokens, int hidden_size) { + const int tid = threadIdx.x; + const int token_idx = blockIdx.x; + + if constexpr (use_per_token_quant) { + float amax_val = 0.0f; + const float zero = 0.0f; + + for (int i = tid; i < hidden_size; i += blockDim.x) { + float val = (float)input[token_idx * hidden_size + i]; + val = val > zero ? val : -val; + if (val > amax_val) + amax_val = val; + } + + __shared__ float s_amax; + const float block_amax_val = blockReduceMax(amax_val); + if (tid == 0) { + s_amax = block_amax_val; + scale[token_idx] = __float2half_rn(block_amax_val / 127.0f); + } + __syncthreads(); + + float tmp_scale = 127.0f / s_amax; + for (int i = tid; i < hidden_size; i += blockDim.x) { + output[token_idx * hidden_size + i] = + float_to_int8_rn(((float)input[token_idx * hidden_size + i]) * tmp_scale); + } + } else { + for (int i = tid; i < hidden_size; i += blockDim.x) { + output[token_idx * hidden_size + i] = + float_to_int8_rn(((float)input[token_idx * hidden_size + i]) / __half2float(scale)); + } + } +} +} + + + +void invoke_quant(torch::Tensor &out, // [..., hidden_size] + torch::Tensor &input, // [..., hidden_size] + torch::Tensor &scale) { // [num_tokens] + assert(input.is_contiguous()); + assert(out.is_contiguous()); + int hidden_size = input.size(-1); + int num_tokens = input.numel() / hidden_size; + dim3 grid(num_tokens); + dim3 block(std::min(hidden_size, 1024)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "quant_kernel", [&] { + vllm::quant_kernel<<>>( + input.data_ptr(), out.data_ptr(), + scale.data_ptr(), num_tokens, hidden_size); + }); +} + diff --git a/llm-awq/awq/kernels/csrc/w8a8/quantization.h b/llm-awq/awq/kernels/csrc/w8a8/quantization.h new file mode 100644 index 0000000000000000000000000000000000000000..a5db72d730aa30075c2548a4cf4b5ac50b7f9b22 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/w8a8/quantization.h @@ -0,0 +1,4 @@ +#include +void invoke_quant(torch::Tensor &out, // [..., hidden_size] + torch::Tensor &input, // [..., hidden_size] + torch::Tensor &scale); // [num_tokens] \ No newline at end of file diff --git a/llm-awq/awq/kernels/csrc/w8a8/reduction_utils.cuh b/llm-awq/awq/kernels/csrc/w8a8/reduction_utils.cuh new file mode 100644 index 0000000000000000000000000000000000000000..205dbe68ac86538ddcced2405fdfe5455b7fa787 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/w8a8/reduction_utils.cuh @@ -0,0 +1,170 @@ +/* + * Adapted from https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/reduce_kernel_utils.cuh + * Copyright (c) 2023, The vLLM team. + * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once +#define FINAL_MASK 0xffffffff + + +namespace vllm { + +template +__inline__ __device__ T warpReduceSum(T val) { +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) + val += __shfl_xor_sync(0xffffffff, val, mask, 32); + return val; +} + +template +__inline__ __device__ T warpReduceSumV2(T* val) +{ +#pragma unroll + for (int i = 0; i < NUM; i++) + { +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) + val[i] += __shfl_xor_sync(FINAL_MASK, val[i], mask, 32); + } + return (T) (0.0f); +} + +/* Calculate the sum of all elements in a block */ +template +__inline__ __device__ T blockReduceSum(T val) { + static __shared__ T shared[32]; + int lane = threadIdx.x & 0x1f; + int wid = threadIdx.x >> 5; + + val = warpReduceSum(val); + + if (lane == 0) + shared[wid] = val; + + __syncthreads(); + + // Modify from blockDim.x << 5 to blockDim.x / 32. to prevent + // blockDim.x is not divided by 32 + val = (threadIdx.x < (blockDim.x / 32.f)) ? shared[lane] : (T)(0.0f); + val = warpReduceSum(val); + return val; +} + +/* Calculate the sum of all elements in a block */ +template +__inline__ __device__ T blockAllReduceSum(T val) { + static __shared__ T shared[32]; + int lane = threadIdx.x & 0x1f; + int wid = threadIdx.x >> 5; + + val = warpReduceSum(val); + + if (lane == 0) + shared[wid] = val; + + __syncthreads(); + + // Modify from blockDim.x << 5 to blockDim.x / 32. to prevent + // blockDim.x is not divided by 32 + val = (lane < (blockDim.x / 32.f)) ? shared[lane] : (T)(0.0f); + val = warpReduceSum(val); + return val; +} + +template +__inline__ __device__ T blockReduceSumV2(T* val) +{ + static __shared__ T shared[NUM][33]; + int lane = threadIdx.x & 0x1f; + int wid = threadIdx.x >> 5; + + warpReduceSumV2(val); + + if (lane == 0) + { +#pragma unroll + for (int i = 0; i < NUM; i++) + { + shared[i][wid] = val[i]; + } + } + + __syncthreads(); + + bool is_mask = threadIdx.x < (blockDim.x / 32.f); +#pragma unroll + for (int i = 0; i < NUM; i++) + { + val[i] = is_mask ? shared[i][lane] : (T) (0.0f); + } + warpReduceSumV2(val); + return (T) 0.0f; +} + +template +__inline__ __device__ T warpReduceMax(T val) +{ +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) + val = max(val, __shfl_xor_sync(0xffffffff, val, mask, 32)); + return val; +} +/* Calculate the maximum of all elements in a block */ +template +__inline__ __device__ T blockReduceMax(T val) +{ + static __shared__ T shared[32]; + int lane = threadIdx.x & 0x1f; // in-warp idx + int wid = threadIdx.x >> 5; // warp idx + val = warpReduceMax(val); // get maxx in each warp + if (lane == 0) // record in-warp maxx by warp Idx + shared[wid] = val; + __syncthreads(); + // Modify from blockDim.x << 5 to blockDim.x / 32. to prevent + // blockDim.x is not divided by 32 + val = (threadIdx.x < (blockDim.x / 32.f)) ? shared[lane] : -1e20f; + val = warpReduceMax(val); + return val; +} + +/* Calculate the maximum of all elements in a block */ +template +__inline__ __device__ T blockAllReduceMax(T val) +{ + static __shared__ T shared[32]; + int lane = threadIdx.x & 0x1f; // in-warp idx + int wid = threadIdx.x >> 5; // warp idx + + val = warpReduceMax(val); // get maxx in each warp + + if (lane == 0) // record in-warp maxx by warp Idx + shared[wid] = val; + + __syncthreads(); + + // Modify from blockDim.x << 5 to blockDim.x / 32. to prevent + // blockDim.x is not divided by 32 + val = (lane < (blockDim.x / 32.f)) ? shared[lane] : -1e20f; + val = warpReduceMax(val); + + return val; +} + + + + + +} // namespace vllm diff --git a/llm-awq/awq/kernels/csrc/w8a8/utils.cuh b/llm-awq/awq/kernels/csrc/w8a8/utils.cuh new file mode 100644 index 0000000000000000000000000000000000000000..cd4d45df253273bb1de6e2652836d86ecee4054b --- /dev/null +++ b/llm-awq/awq/kernels/csrc/w8a8/utils.cuh @@ -0,0 +1,469 @@ +// Adated from FasterTransformer, https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention/decoder_masked_multihead_attention_template.hpp +// Modified by Haotian Tang +#pragma once + +#include +#include +#include +#include +#include + +template +struct FloatTypeConverter +{ + using Type = T; +}; + +template <> +struct FloatTypeConverter +{ + using Type = half; +}; + +template <> +struct FloatTypeConverter +{ + using Type = __nv_bfloat16; +}; + +template <> +struct FloatTypeConverter +{ + using Type = float; +}; + + + +template struct num_elems; +template <> struct num_elems { static constexpr int value = 1; }; +template <> struct num_elems { static constexpr int value = 2; }; +template <> struct num_elems { static constexpr int value = 4; }; +template <> struct num_elems { static constexpr int value = 1; }; +template <> struct num_elems { static constexpr int value = 2; }; +#ifdef ENABLE_BF16 +template <> struct num_elems<__nv_bfloat16> { static constexpr int value = 1; }; +template <> struct num_elems<__nv_bfloat162> { static constexpr int value = 2; }; +#endif +#ifdef ENABLE_FP8 +template <> struct num_elems<__nv_fp8_e4m3> { static constexpr int value = 1; }; +template <> struct num_elems<__nv_fp8x2_e4m3> { static constexpr int value = 2; }; +#endif + +template struct packed_as; +template struct packed_as { using type = T; }; +template<> struct packed_as { using type = half2; }; +template<> struct packed_as { using type = float2; }; +template<> struct packed_as { using type = int16_t; }; +template<> struct packed_as { using type = int2; }; +template<> struct packed_as { using type = half; }; +template<> struct packed_as { using type = float; }; +#ifdef ENABLE_BF16 +template<> struct packed_as<__nv_bfloat16, 2> { using type = __nv_bfloat162; }; +template<> struct packed_as<__nv_bfloat162, 1> { using type = __nv_bfloat16; }; +#endif +#ifdef ENABLE_FP8 +template<> struct packed_as<__nv_fp8_e4m3, 2> { using type = __nv_fp8x2_e4m3; }; +template<> struct packed_as<__nv_fp8x2_e4m3, 1> { using type = __nv_fp8_e4m3; }; +template<> struct packed_as<__nv_fp8_e5m2, 2> { using type = __nv_fp8x2_e5m2; }; +template<> struct packed_as<__nv_fp8x2_e5m2, 1> { using type = __nv_fp8_e5m2; }; +#endif + +inline __device__ float2 operator*(float2 a, float2 b) { return make_float2(a.x * b.x, a.y * b.y); } +inline __device__ float2 operator+(float2 a, float2 b) { return make_float2(a.x + b.x, a.y + b.y); } +inline __device__ float2 operator-(float2 a, float2 b) { return make_float2(a.x - b.x, a.y - b.y); } + +inline __device__ float2 operator*(float2 a, float b) { return make_float2(a.x * b, a.y * b); } +inline __device__ float2 operator+(float2 a, float b) { return make_float2(a.x + b, a.y + b); } +inline __device__ float2 operator-(float2 a, float b) { return make_float2(a.x - b, a.y - b); } + +static inline __device__ int8_t float_to_int8_rn(float x) +{ + uint32_t dst; + asm volatile("cvt.rni.sat.s8.f32 %0, %1;" : "=r"(dst) : "f"(x)); + return reinterpret_cast(dst); +} + +template +inline __device__ T ldg(const T* val) { + return __ldg(val); +} + +#if ENABLE_BF16 +#define bf1622float2 __bfloat1622float2 +#define float22bf162 __float22bfloat162_rn +#define bf162bf162 __bfloat162bfloat162 +inline __device__ int16_t bf1622int16(__nv_bfloat162 val) +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + float2 f_val; + f_val.x = max(min(__low2float(val), 127.f), -128.f); + f_val.y = max(min(__high2float(val), 127.f), -128.f); + + union + { + int8_t int8[2]; + int16_t int16; + }; + + int8[0] = static_cast(static_cast(f_val.x)); + int8[1] = static_cast(static_cast(f_val.y)); + return int16; +#else + val = __hmin2(val, make_bfloat162(127., 127.)); + val = __hmax2(val, make_bfloat162(-128., -128.)); + + union + { + int8_t int8[2]; + int16_t int16; + }; + + int8[0] = static_cast(static_cast(val.x)); + int8[1] = static_cast(static_cast(val.y)); + return int16; +#endif +} +#endif + +#if ENABLE_BF16 +template<> +inline __device__ __nv_bfloat162 ldg(const __nv_bfloat162* val) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + return val[0]; +#else + return __ldg(val); +#endif +} + +template<> +inline __device__ __nv_bfloat16 ldg(const __nv_bfloat16* val) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + return val[0]; +#else + return __ldg(val); +#endif +} +#endif // ENABLE_BF16 + +template +__device__ inline T_OUT cuda_cast(T_IN val) +{ + return val; +} + +template <> +__device__ inline float2 cuda_cast(int2 val) +{ + return make_float2(val.x, val.y); +} + +template <> +__device__ inline float2 cuda_cast(float val) +{ + return make_float2(val, val); +} + +template <> +__device__ inline float2 cuda_cast(half2 val) +{ + return __half22float2(val); +} + +template <> +__device__ inline half2 cuda_cast(float2 val) +{ + return __float22half2_rn(val); +} + +template <> +__device__ inline half2 cuda_cast(float val) +{ + return __float2half2_rn(val); +} + +template <> +__device__ inline half2 cuda_cast(half val) +{ + return __half2half2(val); +} + +template <> +__device__ inline int8_t cuda_cast(half val) +{ + union + { + int8_t int8[2]; + int16_t int16; + }; + + union + { + half fp16; + int16_t int16_in; + }; + + fp16 = val; + asm volatile("cvt.rni.sat.s8.f16 %0, %1;" : "=h"(int16) : "h"(int16_in)); + return int8[0]; +} + +template <> +__device__ inline int16_t cuda_cast(half2 val) +{ + union + { + int8_t int8[2]; + int16_t int16; + }; + + int8[0] = cuda_cast(val.x); + int8[1] = cuda_cast(val.y); + return int16; +} + +template <> +__device__ inline int8_t cuda_cast(float val) +{ + union + { + int8_t int8[2]; + int16_t int16; + }; + + asm volatile("cvt.rni.sat.s8.f32 %0, %1;" : "=h"(int16) : "f"(val)); + return int8[0]; +} + +template <> +__device__ inline int16_t cuda_cast(float2 val) +{ + union + { + int8_t int8[2]; + int16_t int16; + }; + + int8[0] = cuda_cast(val.x); + int8[1] = cuda_cast(val.y); + return int16; +} + +template <> +__device__ inline half2 cuda_cast(int16_t val) +{ + union + { + int8_t int8[2]; + int16_t int16; + }; + + int16 = val; + return make_half2(int8[0], int8[1]); +} + +template <> +__device__ inline float2 cuda_cast(int16_t val) +{ + union + { + int8_t int8[2]; + int16_t int16; + }; + + int16 = val; + return make_float2(int8[0], int8[1]); +} + +#ifdef ENABLE_BF16 +template <> +__device__ inline __nv_bfloat16 cuda_cast(int32_t val) +{ + return static_cast(val); +} + +template <> +__device__ inline __nv_bfloat16 cuda_cast(int8_t val) +{ + return static_cast(val); +} + +template <> +__device__ inline int8_t cuda_cast(__nv_bfloat16 val) +{ + return static_cast(val); +} + +template <> +__device__ inline float cuda_cast(__nv_bfloat16 val) +{ + return __bfloat162float(val); +} + +template <> +__device__ inline float2 cuda_cast(__nv_bfloat162 val) +{ + return bf1622float2(val); +} + +template <> +__device__ inline half cuda_cast(__nv_bfloat16 val) +{ + return __float2half(__bfloat162float(val)); +} + +template <> +__device__ inline int16_t cuda_cast(__nv_bfloat162 val) +{ + return bf1622int16(val); +} + +template <> +__device__ inline __nv_bfloat16 cuda_cast<__nv_bfloat16, float>(float val) +{ + return __float2bfloat16(val); +} + +template <> +__device__ inline __nv_bfloat16 cuda_cast<__nv_bfloat16, half>(half val) +{ + return __float2bfloat16(__half2float(val)); +} + +template <> +__device__ inline __nv_bfloat162 cuda_cast<__nv_bfloat162, __nv_bfloat16>(__nv_bfloat16 val) +{ + return bf162bf162(val); +} + +template <> +__device__ inline __nv_bfloat162 cuda_cast<__nv_bfloat162, float>(float val) +{ + return __float2bfloat162_rn(val); +} + +template <> +__device__ inline __nv_bfloat162 cuda_cast<__nv_bfloat162, float2>(float2 val) +{ + return float22bf162(val); +} + +template <> +__device__ inline __nv_bfloat162 cuda_cast<__nv_bfloat162, int16_t>(int16_t val) +{ + union + { + int8_t int8[2]; + int16_t int16; + }; + + int16 = val; + __nv_bfloat162 res; + res.x = cuda_cast<__nv_bfloat16>(int8[0]); + res.y = cuda_cast<__nv_bfloat16>(int8[1]); + return res; +} + +template <> +__device__ inline __nv_bfloat162 cuda_cast<__nv_bfloat162, half2>(half2 val) +{ + return float22bf162(__half22float2(val)); +} + +#endif // ENABLE BF16 + +template +__device__ inline To cuda_sum(Ti val) +{ + return cuda_cast(val); +}; + +template +__device__ inline To cuda_sum(float2 val) +{ + return cuda_cast(val.x + val.y); +}; + +// Unary maximum: compute the max of a vector type +template +__device__ inline To cuda_max(Ti val) +{ + return cuda_cast(val); +}; + +template <> +__device__ inline float cuda_max(float2 val) +{ + return fmaxf(val.x, val.y); +} + +template <> +__device__ inline half cuda_max(half2 val) +{ + return __hmax(val.x, val.y); +} + +#ifdef ENABLE_BF16 +template <> +__device__ inline __nv_bfloat16 cuda_max(__nv_bfloat162 val) +{ +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)) + return __hmax(val.x, val.y); +#endif +} +#endif + +// Binary maximum: compute the max of two scalar types +template +__device__ inline T cuda_max(T val1, T val2) +{ + return (val1 > val2) ? val1 : val2; +} + +template +__device__ inline T cuda_abs(T val) +{ + assert(false); + return {}; +} + +template <> +__device__ inline float cuda_abs(float val) +{ + return fabs(val); +} + +template <> +__device__ inline float2 cuda_abs(float2 val) +{ + return make_float2(fabs(val.x), fabs(val.y)); +} + +template <> +__device__ inline half cuda_abs(half val) +{ + return __habs(val); +} + +template <> +__device__ inline half2 cuda_abs(half2 val) +{ + return __habs2(val); +} + +#ifdef ENABLE_BF16 + +#if __CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__) +template <> +__device__ inline __nv_bfloat16 cuda_abs(__nv_bfloat16 val) +{ + return __habs(val); +} + +template <> +__device__ inline __nv_bfloat162 cuda_abs(__nv_bfloat162 val) +{ + return __habs2(val); +} +#endif + +#endif // ENABLE_FP16 \ No newline at end of file diff --git a/llm-awq/awq/kernels/csrc/w8a8/w8a8_gemm_cuda.cu b/llm-awq/awq/kernels/csrc/w8a8/w8a8_gemm_cuda.cu new file mode 100644 index 0000000000000000000000000000000000000000..96a72d7683d4d47f05d848e29ac803176ec5c2b6 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/w8a8/w8a8_gemm_cuda.cu @@ -0,0 +1,953 @@ +// Inspired by QServe https://github.com/mit-han-lab/qserve/tree/main. +// Modified by Yuming Lou. +// @article{lin2024awq, +// title={AWQ: Activation-aware Weight Quantization for On-Device LLM Compression and Acceleration}, +// author={Lin, Ji and Tang, Jiaming and Tang, Haotian and Yang, Shang and Chen, Wei-Ming and Wang, Wei-Chen and Xiao, Guangxuan and Dang, Xingyu and Gan, Chuang and Han, Song}, +// journal={Proceedings of Machine Learning and Systems}, +// volume={6}, +// pages={87--100}, +// year={2024} +// } + +#include "w8a8_gemm_cuda.h" +#include +#include +#include + +#define OP_M 16 +#define OP_N 8 +#define OP_K 32 +#define INTRIN_M 16 +#define INTRIN_N 16 +#define INTRIN_K 32 +#define WARP_SIZE 32 +#define SMEM_PAD_A 0 +#define SMEM_PAD_B 0 +#define PACK_SIZE 16 +#if (__CUDACC_VER_MAJOR__ >= 11) && (__CUDACC_VER_MINOR__ >= 4) +#define L2_CACHEHINT(size) ".L2::" #size "B" +#else +#define L2_CACHEHINT(size) +#endif +#define KERNEL_LAUNCH_CODE_FUSE_BIAS \ + constexpr int NUM_WARPS = (CTA_M / WARP_M) * (CTA_N / WARP_N) * (CTA_K / WARP_K); \ + constexpr int kSmemByteSize = \ + (CTA_M * (CTA_K + SMEM_PAD_A) + CTA_N * (CTA_K + SMEM_PAD_B)) * STAGES * \ + sizeof(int8_t) + CTA_N * sizeof(float); \ + if (kSmemByteSize >= 99 * 1024) \ + { \ + printf("This kernel requires %d Bytes of shared memory, which exceeds " \ + "device limit.\n", \ + kSmemByteSize); \ + return ; \ + } \ + int num_blocks_m = (num_out_feats + CTA_M - 1) / CTA_M; \ + int num_blocks_n = (num_out_channels+ CTA_N - 1) / CTA_N / 1; \ + const int log_tile = get_log_tile<8>((num_out_feats + CTA_M - 1) / CTA_M); \ + const int tile_shift = 1 << log_tile; \ + dim3 num_blocks(num_blocks_n *tile_shift, \ + (num_blocks_m + tile_shift - 1) / tile_shift); \ + dim3 threads_per_block(WARP_SIZE, NUM_WARPS); \ + auto kernel_func = \ + dense_kernel0_fuse_bias; \ + cudaFuncSetAttribute(kernel_func, cudaFuncAttributeMaxDynamicSharedMemorySize, \ + kSmemByteSize); \ + kernel_func<<>>( \ + in_feats, kernel, wscales, ascales, out_feats, bias, num_in_feats, num_out_channels, \ + num_in_channels); + + +#define KERNEL_LAUNCH_CODE \ + constexpr int NUM_WARPS = (CTA_M / WARP_M) * (CTA_N / WARP_N) * (CTA_K / WARP_K); \ + constexpr int kSmemByteSize = \ + (CTA_M * (CTA_K + SMEM_PAD_A) + CTA_N * (CTA_K + SMEM_PAD_B)) * STAGES * \ + sizeof(int8_t); \ + if (kSmemByteSize >= 99 * 1024) \ + { \ + printf("This kernel requires %d Bytes of shared memory, which exceeds " \ + "device limit.\n", \ + kSmemByteSize); \ + return ; \ + } \ + int num_blocks_m = (num_out_feats + CTA_M - 1) / CTA_M; \ + int num_blocks_n = num_out_channels / CTA_N / 1; \ + const int log_tile = get_log_tile<8>((num_out_feats + CTA_M - 1) / CTA_M); \ + const int tile_shift = 1 << log_tile; \ + dim3 num_blocks(num_blocks_n *tile_shift, \ + (num_blocks_m + tile_shift - 1) / tile_shift); \ + dim3 threads_per_block(WARP_SIZE, NUM_WARPS); \ + auto kernel_func = \ + dense_kernel0; \ + cudaFuncSetAttribute(kernel_func, cudaFuncAttributeMaxDynamicSharedMemorySize, \ + kSmemByteSize); \ + kernel_func<<>>( \ + in_feats, kernel, wscales, ascales, out_feats, num_in_feats, num_out_channels, \ + num_in_channels); + + + +template +__inline__ __host__ __device__ int get_log_tile(int n) +{ + if (N >= 8 && n >= 6) + return 3; + else if (N >= 4 && n >= 3) + return 2; + else if (N >= 2 && n >= 2) + return 1; + else + return 0; +} + +__inline__ __device__ uint2 get_block_idx_mapping(int blockIdx_x, + int blockIdx_y, + int log_tile) +{ + return make_uint2((blockIdx_x >> log_tile), + (blockIdx_y << log_tile) + + ((blockIdx_x) & ((1 << (log_tile)) - 1))); +} + +__inline__ __device__ uint32_t cast_smem_ptr_to_uint(void const *const ptr) +{ + uint32_t smem_int_ptr; + + asm("{.reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, " + "smem_ptr; }\n" + : "=r"(smem_int_ptr) + : "l"(ptr)); + + return smem_int_ptr; +} + +__inline__ __device__ void ldmatrix_m8n8_x4_b16(int8_t *shared_warp, int ax0_0, + uint32_t addr) +{ + __asm__ __volatile__("ldmatrix.sync.aligned.m8n8.x4.shared.b16" + "{%0, %1, %2, %3}, [%4];" + : "=r"(((unsigned *)(shared_warp + (ax0_0 * 16)))[0]), + "=r"(((unsigned *)(shared_warp + (ax0_0 * 16)))[1]), + "=r"(((unsigned *)(shared_warp + (ax0_0 * 16)))[2]), + "=r"(((unsigned *)(shared_warp + (ax0_0 * 16)))[3]) + : "r"(addr)); +} + +__inline__ __device__ void +ldmatrix_m8n8_x4_trans_b16(int8_t *shared_warp, int ax0_0, uint32_t addr) +{ + __asm__ __volatile__("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16" + "{%0, %1, %2, %3}, [%4];" + : "=r"(((unsigned *)(shared_warp + (ax0_0 * 16)))[0]), + "=r"(((unsigned *)(shared_warp + (ax0_0 * 16)))[1]), + "=r"(((unsigned *)(shared_warp + (ax0_0 * 16)))[2]), + "=r"(((unsigned *)(shared_warp + (ax0_0 * 16)))[3]) + : "r"(addr)); +} + +// function from lmdeploy +__inline__ __device__ void +cp_async_cg_A(uint32_t smem_int_ptr, const uint4 *__restrict__ src, bool mask)//256 * int8 +{ + const int cp_size = 16; + asm volatile("{" + " .reg .pred p;" + " setp.ne.b32 p, %0, 0;" + " @p cp.async.cg.shared.global" L2_CACHEHINT(128) " [%1], [%2], %3;" + "}" ::"r"((int)mask), + "r"(smem_int_ptr), + "l"(src), + "n"(cp_size)); +} + +__device__ __inline__ void mma_m16n8k32(void *C_warp, void *A_shared_warp, + void *B_shared_warp) +{ + __asm__ __volatile__( + "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32" + "{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%10, %11, %12, %13};" + : "=r"(((int *)C_warp)[0]), "=r"(((int *)C_warp)[1]), + "=r"(((int *)C_warp)[2]), "=r"(((int *)C_warp)[3]) + : "r"(((unsigned *)A_shared_warp)[0]), + "r"(((unsigned *)A_shared_warp)[1]), + "r"(((unsigned *)A_shared_warp)[2]), + "r"(((unsigned *)A_shared_warp)[3]), + "r"(((unsigned *)B_shared_warp)[0]), + "r"(((unsigned *)B_shared_warp)[1]), "r"(((int *)C_warp)[0]), + "r"(((int *)C_warp)[1]), "r"(((int *)C_warp)[2]), + "r"(((int *)C_warp)[3])); +} + +template +__device__ __inline__ void +global_to_share_one_stage_A(int8_t *src, int8_t *dst, int global_ncols, + int cta_offset_m, int cta_offset_n, + int global_iter_k, int shared_iter_k, bool mask, + bool *preds) +{ + constexpr int total_global_iters = (CTA_M * CTA_K) / PACK_SIZE / CTA_SIZE; + constexpr int partial_global_iters = total_global_iters / SHARED_K_ITERS; + constexpr int cta_step_m_or_n = (CTA_SIZE * PACK_SIZE) / CTA_K; + constexpr int warp_step_m_or_n = (WARP_SIZE * PACK_SIZE) / CTA_K; + constexpr int threads_per_row = CTA_K / PACK_SIZE; + constexpr int kSmemCol = CTA_K + SMEM_PAD_A; + int8_t *dst_hoisted = dst; + int8_t *src_hoisted = src + global_iter_k * CTA_K; + + if (mask) + { +#pragma unroll + for (int _global_iter = 0; _global_iter < partial_global_iters; + ++_global_iter) + { + int global_iter = shared_iter_k * partial_global_iters + _global_iter; + + void *dst_ptr = + (void *)(dst_hoisted + global_iter * cta_step_m_or_n * kSmemCol); + uint4 *src_ptr = + (uint4 *)(src_hoisted + global_iter * cta_step_m_or_n * global_ncols); + if constexpr (STAGES > 1) + { + uint32_t addr = cast_smem_ptr_to_uint(dst_ptr); + cp_async_cg_A(addr, src_ptr, preds[global_iter]); + } + else + { + if (preds[global_iter]) + *(uint4 *)dst_ptr = *src_ptr; + } + } + } +} + +template +__device__ __inline__ void +global_to_share_one_stage_B(int8_t *src, int8_t *dst, int global_ncols, + int cta_offset_m, int cta_offset_n, + int global_iter_k, int shared_iter_k, bool mask, bool *preds) +{ + constexpr int total_global_iters = (CTA_N * CTA_K) / PACK_SIZE / CTA_SIZE; + constexpr int partial_global_iters = total_global_iters / SHARED_K_ITERS; + constexpr int cta_step_m_or_n = (CTA_SIZE * PACK_SIZE) / CTA_K; + constexpr int warp_step_m_or_n = (WARP_SIZE * PACK_SIZE) / CTA_K; + constexpr int threads_per_row = CTA_K / PACK_SIZE; + constexpr int kSmemCol = CTA_K + SMEM_PAD_B; + int8_t *dst_hoisted = dst; + int8_t *src_hoisted = src + global_iter_k * CTA_K; +#pragma unroll + for (int _global_iter = 0; _global_iter < partial_global_iters; + ++_global_iter) + { + int global_iter = shared_iter_k * partial_global_iters + _global_iter; + + void *dst_ptr = + (void *)(dst_hoisted + global_iter * cta_step_m_or_n * kSmemCol); + uint4 *src_ptr = + (uint4 *)(src_hoisted + global_iter * cta_step_m_or_n * global_ncols); + if constexpr (STAGES > 1) + { + uint32_t addr = cast_smem_ptr_to_uint(dst_ptr); + cp_async_cg_A(addr, src_ptr, preds[global_iter]); + } + else + { + if (preds[global_iter]) + *(uint4 *)dst_ptr = *src_ptr; + } + } +} + +template +__device__ __inline__ void +share_to_reg_one_stage_A(int8_t *src, int8_t *dst, int warp_offset_m, + int warp_offset_n, int k_0_1, int shared_iters) +{ + constexpr int kSmemCol = CTA_K + SMEM_PAD_A; + int ld_col = (k_0_1 * INTRIN_K + (threadIdx.x / 16) * 16) / PACK_SIZE; + + for (int shared_iter = 0; shared_iter < shared_iters; ++shared_iter) + { + int ld_row = warp_offset_m + shared_iter * INTRIN_M + (threadIdx.x % 16); + int ld_col_swizzled = ld_col ^ (ld_row / 2) & 3; + void *addr_ptr = + (void *)(src + ld_row * kSmemCol + ld_col_swizzled * PACK_SIZE); + uint32_t addr = cast_smem_ptr_to_uint(addr_ptr); + ldmatrix_m8n8_x4_b16(dst, shared_iter, addr); + } +} + +template +__device__ __inline__ void +share_to_reg_one_stage_B(int8_t *src, int8_t *dst, int warp_offset_m, + int warp_offset_n, int k_0_1, int shared_iters) +{ + constexpr int kSmemCol = CTA_K + SMEM_PAD_B; + int ld_col = (k_0_1 * INTRIN_K + ((threadIdx.x / 8) % 2) * 16) / PACK_SIZE; + + for (int shared_iter = 0; shared_iter < shared_iters; ++shared_iter) + { + int ld_row = warp_offset_n + shared_iter * INTRIN_N + ((threadIdx.x / 8 / 2) * 8 + threadIdx.x % 8); + int ld_col_swizzled = ld_col ^ (ld_row / 2) & 3; + void *addr_ptr = + (void *)(src + ld_row * kSmemCol + ld_col_swizzled * PACK_SIZE); + uint32_t addr = cast_smem_ptr_to_uint(addr_ptr); + ldmatrix_m8n8_x4_b16(dst, shared_iter, addr); + } +} + +template +__global__ void dense_kernel0_fuse_bias(int8_t *__restrict__ A, int8_t *__restrict__ B, + half2 *__restrict__ wscales, half *__restrict__ ascales, + half *__restrict__ C, half *__restrict__ Bias, + int M, int N, int K) +{ + constexpr int NUM_WARPS_MN = CTA_M / WARP_M * CTA_N / WARP_N; + constexpr int NUM_WARPS = NUM_WARPS_MN * CTA_K / WARP_K; + constexpr int CTA_SIZE = NUM_WARPS * WARP_SIZE; + constexpr int CTA_SIZE_MN = NUM_WARPS_MN * WARP_SIZE; + constexpr int SLICES = CTA_K / WARP_K; + int num_blocks_n = (N + CTA_N - 1) / CTA_N; + int num_blocks_m = (M + CTA_M - 1) / CTA_M; + + int blockIdx_n = blockIdx.x; + int blockIdx_m = blockIdx.y; + const int log_tile = get_log_tile<8>((M + CTA_M - 1) / CTA_M); + const uint2 block_idx_mapping = + get_block_idx_mapping(blockIdx_n, blockIdx_m, log_tile); + blockIdx_n = block_idx_mapping.x; + blockIdx_m = block_idx_mapping.y; + + int C_warp[CTA_M * CTA_N / CTA_SIZE_MN]; + constexpr int kSmemPadKA = CTA_K + SMEM_PAD_A; + constexpr int kSmemPadKB = CTA_K + SMEM_PAD_B; + constexpr int kSmemSizeAPerStage = CTA_M * kSmemPadKA; + constexpr int kSmemSizeBPerStage = CTA_N * kSmemPadKB; + constexpr int kSmemSizeA = kSmemSizeAPerStage * STAGES; + constexpr int kSmemSizeB = kSmemSizeBPerStage * STAGES; + extern __shared__ int8_t mem_shared[]; + int8_t *A_shared = mem_shared; + int8_t *B_shared = mem_shared + kSmemSizeA; + float *Bias_shared= reinterpret_cast(mem_shared + kSmemSizeA + kSmemSizeB); + int8_t A_shared_warp_[2][WARP_M * WARP_K / + WARP_SIZE]; + int8_t B_shared_warp_[2][WARP_N * WARP_K / + WARP_SIZE]; + constexpr int A_total_global_iters = (CTA_M * CTA_K) / PACK_SIZE / CTA_SIZE; + constexpr int B_total_global_iters = (CTA_N * CTA_K) / PACK_SIZE / CTA_SIZE; + constexpr int A_src_step_m = (CTA_SIZE * PACK_SIZE) / CTA_K; + constexpr int B_src_step_k = (CTA_SIZE * PACK_SIZE) / CTA_K; + constexpr int A_warp_step_m = (WARP_SIZE * PACK_SIZE) / CTA_K; + constexpr int B_warp_step_n = (WARP_SIZE * PACK_SIZE) / CTA_K; + constexpr int A_threads_per_row = CTA_K / PACK_SIZE; + constexpr int B_threads_per_row = CTA_K / PACK_SIZE; + int cta_offset_m = blockIdx_m * CTA_M; + int cta_offset_n = blockIdx_n * CTA_N; + int warp_mn = threadIdx.y % NUM_WARPS_MN; + int slice_id = threadIdx.y / NUM_WARPS_MN; // Always zero if threadIdx.z==0! + int warp_offset_m = (warp_mn % (CTA_M / WARP_M)) * WARP_M; + int warp_offset_n = (warp_mn / (CTA_M / WARP_M)) * WARP_N; + int warp_offset_k = slice_id * WARP_K; + + for (int i = 0; i < CTA_M * CTA_N / CTA_SIZE_MN; i++) + C_warp[i] = 0; + + int gemm_iters = (K + CTA_K - 1) / CTA_K; + int k_0_0_ld = 0; + int k_0_0 = 0; + constexpr int prologue_stages = STAGES == 1 ? 1 : STAGES - 1; + int A_hoisted_row = threadIdx.y * A_warp_step_m + (threadIdx.x / A_threads_per_row); + int A_hoisted_col = (threadIdx.x % A_threads_per_row); + int A_hoisted_col_swizzled = A_hoisted_col ^ (A_hoisted_row / 2) & 3; + + int B_hoisted_row = threadIdx.y * B_warp_step_n + (threadIdx.x / B_threads_per_row); + int B_hoisted_col = (threadIdx.x % B_threads_per_row); + int B_hoisted_col_swizzled = B_hoisted_col ^ (B_hoisted_row / 2) & 3; + + int8_t *A_shared_hoisted = A_shared + + A_hoisted_row * kSmemPadKA + + A_hoisted_col_swizzled * PACK_SIZE; + int8_t *B_shared_hoisted = B_shared + B_hoisted_row * kSmemPadKB + + B_hoisted_col_swizzled * PACK_SIZE; + int8_t *A_hoisted = A + cta_offset_m * K + A_hoisted_row * K + + A_hoisted_col * PACK_SIZE; + int8_t *B_hoisted = B + cta_offset_n * K + B_hoisted_row * K + + B_hoisted_col * PACK_SIZE; + bool A_g2s_preds[A_total_global_iters]; + bool B_g2s_preds[B_total_global_iters]; + //debug + // printf("A: %d ",A_total_global_iters); + // printf("B: %d ",B_total_global_iters); + // printf("prologue_stages: %d ",prologue_stages); + // __shared__ float2 Bias_shared[CTA_N]; + #pragma unroll + for (int i = 0; i < CTA_N ; i++) + { + Bias_shared[i] = __half2float(Bias[cta_offset_n+i]); + } + + +#pragma unroll + for (int i = 0; i < A_total_global_iters; i++) + { + A_g2s_preds[i] = (cta_offset_m + A_hoisted_row + i * A_src_step_m) < M; + } + #pragma unroll + for (int i = 0; i < B_total_global_iters; i++) + { + B_g2s_preds[i] = cta_offset_n + B_hoisted_row + i * B_src_step_k < N; + } + int *C_shared = reinterpret_cast(mem_shared); +#pragma unroll + for (k_0_0_ld = 0; k_0_0_ld < prologue_stages; ++k_0_0_ld) + { + global_to_share_one_stage_A( + A_hoisted, A_shared_hoisted + k_0_0_ld * kSmemSizeAPerStage, K, + cta_offset_m, cta_offset_n, k_0_0_ld, 0, true, A_g2s_preds); + global_to_share_one_stage_B( + B_hoisted, B_shared_hoisted + k_0_0_ld * kSmemSizeBPerStage, K, + cta_offset_m, cta_offset_n, k_0_0_ld, 0, true, B_g2s_preds); + if constexpr (STAGES > 1) + __pipeline_commit(); + } + if constexpr (STAGES > 1) + __pipeline_wait_prior(STAGES - 2); + __syncthreads(); + +// global_to_share_bias(Bias,Bias_shared,cta_offset_n); + + share_to_reg_one_stage_A( + A_shared + warp_offset_k, A_shared_warp_[0], warp_offset_m, warp_offset_n, 0, + WARP_M / INTRIN_M); + share_to_reg_one_stage_B( + B_shared + warp_offset_k, B_shared_warp_[0], warp_offset_m, warp_offset_n, 0, + WARP_N / INTRIN_N); + constexpr int SHARED_K_ITERS = WARP_K / INTRIN_K; + + for (; k_0_0 < gemm_iters; ++k_0_0, ++k_0_0_ld) + { + int ld_stage = k_0_0_ld % STAGES; + int compute_stage = k_0_0 % STAGES; + int8_t *A_shared_this_compute_stage; + int8_t *B_shared_this_compute_stage; + + for (int iter_k = 0; iter_k < SHARED_K_ITERS; ++iter_k) + { + A_shared_this_compute_stage = + A_shared + compute_stage * kSmemSizeAPerStage + warp_offset_k; + B_shared_this_compute_stage = + B_shared + compute_stage * kSmemSizeBPerStage + warp_offset_k; + share_to_reg_one_stage_A( + A_shared_this_compute_stage, A_shared_warp_[(iter_k + 1) % 2], + warp_offset_m, warp_offset_n, (iter_k + 1) % SHARED_K_ITERS, + WARP_M / INTRIN_M); + share_to_reg_one_stage_B( + B_shared_this_compute_stage, B_shared_warp_[(iter_k + 1) % 2], + warp_offset_m, warp_offset_n, (iter_k + 1) % SHARED_K_ITERS, + WARP_N / INTRIN_N); + int8_t *A_shared_warp = A_shared_warp_[iter_k % 2]; + int8_t *B_shared_warp = B_shared_warp_[iter_k % 2]; + for (int i_0_3 = 0; i_0_3 < WARP_M / INTRIN_M; ++i_0_3) + { + for (int j_0_4 = 0; j_0_4 < WARP_N / INTRIN_N; ++j_0_4) + { + mma_m16n8k32( + (void *)(C_warp + i_0_3 * WARP_N / INTRIN_N * 8 + j_0_4 * 8), + (void *)(A_shared_warp + i_0_3 * 16), + (void *)(B_shared_warp + j_0_4 * 16)); + mma_m16n8k32( + (void *)(C_warp + i_0_3 * WARP_N / INTRIN_N * 8 + j_0_4 * 8 + 4), + (void *)(A_shared_warp + i_0_3 * 16), + (void *)(B_shared_warp + j_0_4 * 16 + 8)); + } + } + + if (iter_k < SHARED_K_ITERS - 1) + { + if constexpr (STAGES == 1) + __syncthreads(); + global_to_share_one_stage_A( + A_hoisted, A_shared_hoisted + ld_stage * kSmemSizeAPerStage, K, + cta_offset_m, cta_offset_n, k_0_0_ld, iter_k, + k_0_0_ld < gemm_iters, A_g2s_preds); + global_to_share_one_stage_B( + B_hoisted, B_shared_hoisted + ld_stage * kSmemSizeBPerStage, K, + cta_offset_m, cta_offset_n, k_0_0_ld, iter_k, + k_0_0_ld < gemm_iters, B_g2s_preds); + } + + if (iter_k == SHARED_K_ITERS - 2) + { + if constexpr (STAGES == 1 && SHARED_K_ITERS > 2) + { + __syncthreads(); + } + global_to_share_one_stage_A( + A_hoisted, A_shared_hoisted + ld_stage * kSmemSizeAPerStage, K, + cta_offset_m, cta_offset_n, k_0_0_ld, iter_k + 1, + k_0_0_ld < gemm_iters, A_g2s_preds); + global_to_share_one_stage_B( + B_hoisted, B_shared_hoisted + ld_stage * kSmemSizeBPerStage, K, + cta_offset_m, cta_offset_n, k_0_0_ld, iter_k + 1, + k_0_0_ld < gemm_iters, B_g2s_preds); + if constexpr (STAGES > 1) + { + __pipeline_commit(); + __pipeline_wait_prior(STAGES - 2); + } + compute_stage = (k_0_0 + 1) % STAGES; + __syncthreads(); + } + } + } + __pipeline_commit(); + __pipeline_wait_prior(0); + __syncthreads(); + + if constexpr (SLICES > 1) + { +#pragma unroll + for (int z = 0; z < SLICES; ++z) + { + if (slice_id == z) + { +#pragma unroll + for (int ax0_0_1 = 0; ax0_0_1 < WARP_M / INTRIN_M; ++ax0_0_1) + { +#pragma unroll + for (int ax1_0_1 = 0; ax1_0_1 < WARP_N / INTRIN_N; ++ax1_0_1) + { +#pragma unroll + for (int local_id = 0; local_id < OP_M * 16 / WARP_SIZE; ++local_id) + { + if (z > 0) + { + C_warp[ax0_0_1 * WARP_N / INTRIN_N * 8 + ax1_0_1 * 8 + local_id] += C_shared[warp_offset_m * CTA_N + ax0_0_1 * OP_M * CTA_N + warp_offset_n + ax1_0_1 * 16 + ((local_id % 4) / 2 * 8 + (threadIdx.x / 4)) * CTA_N + (local_id / 4) * 8 + (local_id % 2) + (threadIdx.x % 4) * 2]; + } + C_shared[warp_offset_m * CTA_N + ax0_0_1 * OP_M * CTA_N + warp_offset_n + ax1_0_1 * 16 + ((local_id % 4) / 2 * 8 + (threadIdx.x / 4)) * CTA_N + (local_id / 4) * 8 + (local_id % 2) + (threadIdx.x % 4) * 2] = C_warp[ax0_0_1 * WARP_N / INTRIN_N * 8 + ax1_0_1 * 8 + local_id]; + }; + } + } + } + __syncthreads(); + } + if (slice_id == 0) + { +#pragma unroll + for (int ax0_0_1 = 0; ax0_0_1 < WARP_M / INTRIN_M; ++ax0_0_1) + { +#pragma unroll + for (int ax1_0_1 = 0; ax1_0_1 < WARP_N / INTRIN_N; ++ax1_0_1) + { +#pragma unroll + for (int local_id = 0; local_id < OP_M * 16 / WARP_SIZE; ++local_id) + { + C_warp[ax0_0_1 * WARP_N / INTRIN_N * 8 + ax1_0_1 * 8 + local_id] = C_shared[warp_offset_m * CTA_N + ax0_0_1 * OP_M * CTA_N + warp_offset_n + ax1_0_1 * 16 + ((local_id % 4) / 2 * 8 + (threadIdx.x / 4)) * CTA_N + (local_id / 4) * 8 + (local_id % 2) + (threadIdx.x % 4) * 2]; + }; + } + } + } + } + + int row_wb_thd = cta_offset_m + warp_offset_m + (threadIdx.x / 4); + int col_wb_thd = cta_offset_n + warp_offset_n + (threadIdx.x % 4) * 2; + if (slice_id == 0) + { + for (int ax0_0_1 = 0; ax0_0_1 < WARP_M / INTRIN_M; ++ax0_0_1) + { + int row_wb_1 = row_wb_thd + ax0_0_1 * OP_M; + for (int ax1_0_1 = 0; ax1_0_1 < WARP_N / INTRIN_N; ++ax1_0_1) + { + int col_wb_1 = col_wb_thd + ax1_0_1 * 16; + int *C_warp_local = C_warp + ax0_0_1 * WARP_N / INTRIN_N * 8 + ax1_0_1 * 8; + for (int local_id = 0; local_id < OP_M * 16 / WARP_SIZE; local_id += 2) + { + int row_wb = row_wb_1 + (local_id % 4) / 2 * 8; + int col_wb = col_wb_1 + (local_id / 4) * 8 + (local_id % 2); + if (row_wb < M && col_wb < N ){ + float2 wscale = __half22float2(*(wscales + col_wb / 2)); + float ascale = __half2float(ascales[row_wb]); + float2 psums = make_float2(__int2float_rn(C_warp_local[local_id]), __int2float_rn(C_warp_local[local_id + 1])); + psums.x = psums.x * wscale.x * ascale + Bias_shared[col_wb % CTA_N]; + psums.y = psums.y * wscale.y * ascale + Bias_shared[col_wb % CTA_N + 1]; + *reinterpret_cast(C + row_wb * N + col_wb) = __float22half2_rn(psums); + } + }; + } + } + } +} + +void w8a8_gemm_fuse_bias_forward_cuda(torch::Tensor _in_feats, + torch::Tensor _kernel, + torch::Tensor _wscales, + torch::Tensor _ascales, + torch::Tensor _out_feats, + torch::Tensor _bias) +{ + int num_in_feats = _in_feats.size(0); + int num_in_channels = _in_feats.size(1); + auto in_feats = reinterpret_cast(_in_feats.data_ptr()); + auto kernel = reinterpret_cast(_kernel.data_ptr()); + auto wscales = reinterpret_cast(_wscales.data_ptr()); + auto ascales = reinterpret_cast(_ascales.data_ptr()); + auto bias = reinterpret_cast(_bias.data_ptr()); + // auto options = + // torch::TensorOptions().dtype(torch::kFloat16).device(_in_feats.device()); + // at::Tensor _out_feats = + // torch::empty({num_in_feats, _kernel.size(0)}, options); + int num_out_feats = _out_feats.size(-2); + int num_out_channels = _out_feats.size(-1); + + + auto out_feats = reinterpret_cast(_out_feats.data_ptr()); + + if (num_out_feats > 128) + { + constexpr int CTA_M = 128; + constexpr int CTA_N = 128; + constexpr int CTA_K = 64; + constexpr int WARP_M = 64; + constexpr int WARP_N = 32; + constexpr int WARP_K = 64; + constexpr int STAGES = 6; + KERNEL_LAUNCH_CODE_FUSE_BIAS + } + else + { + constexpr int CTA_M = 64; + constexpr int CTA_N = 64; + constexpr int CTA_K = 64; + constexpr int WARP_M = 32; + constexpr int WARP_N = 16; + constexpr int WARP_K = 64; + constexpr int STAGES = 6; + KERNEL_LAUNCH_CODE_FUSE_BIAS + } + return ; +} + +template +__global__ void dense_kernel0(int8_t *__restrict__ A, int8_t *__restrict__ B, + half2 *__restrict__ wscales, half *__restrict__ ascales, + half *__restrict__ C, int M, int N, int K) +{ + constexpr int NUM_WARPS_MN = CTA_M / WARP_M * CTA_N / WARP_N; + constexpr int NUM_WARPS = NUM_WARPS_MN * CTA_K / WARP_K; + constexpr int CTA_SIZE = NUM_WARPS * WARP_SIZE; + constexpr int CTA_SIZE_MN = NUM_WARPS_MN * WARP_SIZE; + constexpr int SLICES = CTA_K / WARP_K; + int num_blocks_n = (N + CTA_N - 1) / CTA_N; + int num_blocks_m = (M + CTA_M - 1) / CTA_M; + + int blockIdx_n = blockIdx.x; + int blockIdx_m = blockIdx.y; + const int log_tile = get_log_tile<8>((M + CTA_M - 1) / CTA_M); + const uint2 block_idx_mapping = + get_block_idx_mapping(blockIdx_n, blockIdx_m, log_tile); + blockIdx_n = block_idx_mapping.x; + blockIdx_m = block_idx_mapping.y; + + int C_warp[CTA_M * CTA_N / CTA_SIZE_MN]; + constexpr int kSmemPadKA = CTA_K + SMEM_PAD_A; + constexpr int kSmemPadKB = CTA_K + SMEM_PAD_B; + constexpr int kSmemSizeAPerStage = CTA_M * kSmemPadKA; + constexpr int kSmemSizeBPerStage = CTA_N * kSmemPadKB; + constexpr int kSmemSizeA = kSmemSizeAPerStage * STAGES; + constexpr int kSmemSizeB = kSmemSizeBPerStage * STAGES; + extern __shared__ int8_t mem_shared[]; + int8_t *A_shared = mem_shared; + int8_t *B_shared = mem_shared + kSmemSizeA; + int8_t A_shared_warp_[2][WARP_M * WARP_K / + WARP_SIZE]; + int8_t B_shared_warp_[2][WARP_N * WARP_K / + WARP_SIZE]; + constexpr int A_total_global_iters = (CTA_M * CTA_K) / PACK_SIZE / CTA_SIZE; + constexpr int B_total_global_iters = (CTA_N * CTA_K) / PACK_SIZE / CTA_SIZE; + constexpr int A_src_step_m = (CTA_SIZE * PACK_SIZE) / CTA_K; + constexpr int B_src_step_k = (CTA_SIZE * PACK_SIZE) / CTA_K; + constexpr int A_warp_step_m = (WARP_SIZE * PACK_SIZE) / CTA_K; + constexpr int B_warp_step_n = (WARP_SIZE * PACK_SIZE) / CTA_K; + constexpr int A_threads_per_row = CTA_K / PACK_SIZE; + constexpr int B_threads_per_row = CTA_K / PACK_SIZE; + int cta_offset_m = blockIdx_m * CTA_M; + int cta_offset_n = blockIdx_n * CTA_N; + int warp_mn = threadIdx.y % NUM_WARPS_MN; + int slice_id = threadIdx.y / NUM_WARPS_MN; + int warp_offset_m = (warp_mn % (CTA_M / WARP_M)) * WARP_M; + int warp_offset_n = (warp_mn / (CTA_M / WARP_M)) * WARP_N; + int warp_offset_k = slice_id * WARP_K; + + for (int i = 0; i < CTA_M * CTA_N / CTA_SIZE_MN; i++) + C_warp[i] = 0; + + int gemm_iters = (K + CTA_K - 1) / CTA_K; + int k_0_0_ld = 0; + int k_0_0 = 0; + constexpr int prologue_stages = STAGES == 1 ? 1 : STAGES - 1; + int A_hoisted_row = threadIdx.y * A_warp_step_m + (threadIdx.x / A_threads_per_row); + int A_hoisted_col = (threadIdx.x % A_threads_per_row); + int A_hoisted_col_swizzled = A_hoisted_col ^ (A_hoisted_row / 2) & 3; + + int B_hoisted_row = threadIdx.y * B_warp_step_n + (threadIdx.x / B_threads_per_row); + int B_hoisted_col = (threadIdx.x % B_threads_per_row); + int B_hoisted_col_swizzled = B_hoisted_col ^ (B_hoisted_row / 2) & 3; + + int8_t *A_shared_hoisted = A_shared + + A_hoisted_row * kSmemPadKA + + A_hoisted_col_swizzled * PACK_SIZE; + int8_t *B_shared_hoisted = B_shared + B_hoisted_row * kSmemPadKB + + B_hoisted_col_swizzled * PACK_SIZE; + int8_t *A_hoisted = A + cta_offset_m * K + A_hoisted_row * K + + A_hoisted_col * PACK_SIZE; + int8_t *B_hoisted = B + cta_offset_n * K + B_hoisted_row * K + + B_hoisted_col * PACK_SIZE; + bool A_g2s_preds[A_total_global_iters]; +#pragma unroll + for (int i = 0; i < A_total_global_iters; i++) + { + A_g2s_preds[i] = (cta_offset_m + A_hoisted_row + i * A_src_step_m) < M; + } + bool B_g2s_preds[B_total_global_iters]; + #pragma unroll + for (int i = 0; i < B_total_global_iters; i++) + { + B_g2s_preds[i] = (cta_offset_n + B_hoisted_col + i) < N; + } + int *C_shared = reinterpret_cast(mem_shared); +#pragma unroll + for (k_0_0_ld = 0; k_0_0_ld < prologue_stages; ++k_0_0_ld) + { + global_to_share_one_stage_A( + A_hoisted, A_shared_hoisted + k_0_0_ld * kSmemSizeAPerStage, K, + cta_offset_m, cta_offset_n, k_0_0_ld, 0, true, A_g2s_preds); + global_to_share_one_stage_B( + B_hoisted, B_shared_hoisted + k_0_0_ld * kSmemSizeBPerStage, K, + cta_offset_m, cta_offset_n, k_0_0_ld, 0, true, B_g2s_preds); + if constexpr (STAGES > 1) + __pipeline_commit(); + } + if constexpr (STAGES > 1) + __pipeline_wait_prior(STAGES - 2); + __syncthreads(); + + share_to_reg_one_stage_A( + A_shared + warp_offset_k, A_shared_warp_[0], warp_offset_m, warp_offset_n, 0, + WARP_M / INTRIN_M); + share_to_reg_one_stage_B( + B_shared + warp_offset_k, B_shared_warp_[0], warp_offset_m, warp_offset_n, 0, + WARP_N / INTRIN_N); + constexpr int SHARED_K_ITERS = WARP_K / INTRIN_K; + + for (; k_0_0 < gemm_iters; ++k_0_0, ++k_0_0_ld) + { + int ld_stage = k_0_0_ld % STAGES; + int compute_stage = k_0_0 % STAGES; + int8_t *A_shared_this_compute_stage; + int8_t *B_shared_this_compute_stage; + + for (int iter_k = 0; iter_k < SHARED_K_ITERS; ++iter_k) + { + A_shared_this_compute_stage = + A_shared + compute_stage * kSmemSizeAPerStage + warp_offset_k; + B_shared_this_compute_stage = + B_shared + compute_stage * kSmemSizeBPerStage + warp_offset_k; + share_to_reg_one_stage_A( + A_shared_this_compute_stage, A_shared_warp_[(iter_k + 1) % 2], + warp_offset_m, warp_offset_n, (iter_k + 1) % SHARED_K_ITERS, + WARP_M / INTRIN_M); + share_to_reg_one_stage_B( + B_shared_this_compute_stage, B_shared_warp_[(iter_k + 1) % 2], + warp_offset_m, warp_offset_n, (iter_k + 1) % SHARED_K_ITERS, + WARP_N / INTRIN_N); + int8_t *A_shared_warp = A_shared_warp_[iter_k % 2]; + int8_t *B_shared_warp = B_shared_warp_[iter_k % 2]; + for (int i_0_3 = 0; i_0_3 < WARP_M / INTRIN_M; ++i_0_3) + { + for (int j_0_4 = 0; j_0_4 < WARP_N / INTRIN_N; ++j_0_4) + { + mma_m16n8k32( + (void *)(C_warp + i_0_3 * WARP_N / INTRIN_N * 8 + j_0_4 * 8), + (void *)(A_shared_warp + i_0_3 * 16), + (void *)(B_shared_warp + j_0_4 * 16)); + mma_m16n8k32( + (void *)(C_warp + i_0_3 * WARP_N / INTRIN_N * 8 + j_0_4 * 8 + 4), + (void *)(A_shared_warp + i_0_3 * 16), + (void *)(B_shared_warp + j_0_4 * 16 + 8)); + } + } + + if (iter_k < SHARED_K_ITERS - 1) + { + if constexpr (STAGES == 1) + __syncthreads(); + global_to_share_one_stage_A( + A_hoisted, A_shared_hoisted + ld_stage * kSmemSizeAPerStage, K, + cta_offset_m, cta_offset_n, k_0_0_ld, iter_k, + k_0_0_ld < gemm_iters, A_g2s_preds); + global_to_share_one_stage_B( + B_hoisted, B_shared_hoisted + ld_stage * kSmemSizeBPerStage, K, + cta_offset_m, cta_offset_n, k_0_0_ld, iter_k, + k_0_0_ld < gemm_iters, B_g2s_preds); + } + + if (iter_k == SHARED_K_ITERS - 2) + { + if constexpr (STAGES == 1 && SHARED_K_ITERS > 2) + { + __syncthreads(); + } + global_to_share_one_stage_A( + A_hoisted, A_shared_hoisted + ld_stage * kSmemSizeAPerStage, K, + cta_offset_m, cta_offset_n, k_0_0_ld, iter_k + 1, + k_0_0_ld < gemm_iters, A_g2s_preds); + global_to_share_one_stage_B( + B_hoisted, B_shared_hoisted + ld_stage * kSmemSizeBPerStage, K, + cta_offset_m, cta_offset_n, k_0_0_ld, iter_k + 1, + k_0_0_ld < gemm_iters, B_g2s_preds); + if constexpr (STAGES > 1) + { + __pipeline_commit(); + __pipeline_wait_prior(STAGES - 2); + } + compute_stage = (k_0_0 + 1) % STAGES; + __syncthreads(); + } + } + } + __pipeline_commit(); + __pipeline_wait_prior(0); + __syncthreads(); + + if constexpr (SLICES > 1) + { +#pragma unroll + for (int z = 0; z < SLICES; ++z) + { + if (slice_id == z) + { +#pragma unroll + for (int ax0_0_1 = 0; ax0_0_1 < WARP_M / INTRIN_M; ++ax0_0_1) + { +#pragma unroll + for (int ax1_0_1 = 0; ax1_0_1 < WARP_N / INTRIN_N; ++ax1_0_1) + { +#pragma unroll + for (int local_id = 0; local_id < OP_M * 16 / WARP_SIZE; ++local_id) + { + if (z > 0) + { + C_warp[ax0_0_1 * WARP_N / INTRIN_N * 8 + ax1_0_1 * 8 + local_id] += C_shared[warp_offset_m * CTA_N + ax0_0_1 * OP_M * CTA_N + warp_offset_n + ax1_0_1 * 16 + ((local_id % 4) / 2 * 8 + (threadIdx.x / 4)) * CTA_N + (local_id / 4) * 8 + (local_id % 2) + (threadIdx.x % 4) * 2]; + } + C_shared[warp_offset_m * CTA_N + ax0_0_1 * OP_M * CTA_N + warp_offset_n + ax1_0_1 * 16 + ((local_id % 4) / 2 * 8 + (threadIdx.x / 4)) * CTA_N + (local_id / 4) * 8 + (local_id % 2) + (threadIdx.x % 4) * 2] = C_warp[ax0_0_1 * WARP_N / INTRIN_N * 8 + ax1_0_1 * 8 + local_id]; + }; + } + } + } + __syncthreads(); + } + if (slice_id == 0) + { +#pragma unroll + for (int ax0_0_1 = 0; ax0_0_1 < WARP_M / INTRIN_M; ++ax0_0_1) + { +#pragma unroll + for (int ax1_0_1 = 0; ax1_0_1 < WARP_N / INTRIN_N; ++ax1_0_1) + { +#pragma unroll + for (int local_id = 0; local_id < OP_M * 16 / WARP_SIZE; ++local_id) + { + C_warp[ax0_0_1 * WARP_N / INTRIN_N * 8 + ax1_0_1 * 8 + local_id] = C_shared[warp_offset_m * CTA_N + ax0_0_1 * OP_M * CTA_N + warp_offset_n + ax1_0_1 * 16 + ((local_id % 4) / 2 * 8 + (threadIdx.x / 4)) * CTA_N + (local_id / 4) * 8 + (local_id % 2) + (threadIdx.x % 4) * 2]; + }; + } + } + } + } + + int row_wb_thd = cta_offset_m + warp_offset_m + (threadIdx.x / 4); + int col_wb_thd = cta_offset_n + warp_offset_n + (threadIdx.x % 4) * 2; + if (slice_id == 0) + { + for (int ax0_0_1 = 0; ax0_0_1 < WARP_M / INTRIN_M; ++ax0_0_1) + { + int row_wb_1 = row_wb_thd + ax0_0_1 * OP_M; + for (int ax1_0_1 = 0; ax1_0_1 < WARP_N / INTRIN_N; ++ax1_0_1) + { + int col_wb_1 = col_wb_thd + ax1_0_1 * 16; + int *C_warp_local = C_warp + ax0_0_1 * WARP_N / INTRIN_N * 8 + ax1_0_1 * 8; + for (int local_id = 0; local_id < OP_M * 16 / WARP_SIZE; local_id += 2) + { + int row_wb = row_wb_1 + (local_id % 4) / 2 * 8; + int col_wb = col_wb_1 + (local_id / 4) * 8 + (local_id % 2); + if (row_wb < M && col_wb < N){ + int col_wb = col_wb_1 + (local_id / 4) * 8 + (local_id % 2); + float2 wscale = __half22float2(*(wscales + col_wb / 2)); + float ascale = __half2float(ascales[row_wb]); + float2 psums = make_float2(__int2float_rn(C_warp_local[local_id]), __int2float_rn(C_warp_local[local_id + 1])); + psums.x *= wscale.x * ascale; + psums.y *= wscale.y * ascale; + *reinterpret_cast(C + row_wb * N + col_wb) = __float22half2_rn(psums); + } + }; + } + } + } +} + +void w8a8_gemm_forward_cuda(torch::Tensor _in_feats, + torch::Tensor _kernel, + torch::Tensor _wscales, + torch::Tensor _ascales, + torch::Tensor _out_feats) +{ + int num_in_feats = _in_feats.size(0); + int num_in_channels = _in_feats.size(1); + auto in_feats = reinterpret_cast(_in_feats.data_ptr()); + auto kernel = reinterpret_cast(_kernel.data_ptr()); + auto wscales = reinterpret_cast(_wscales.data_ptr()); + auto ascales = reinterpret_cast(_ascales.data_ptr()); + + // auto options = + // torch::TensorOptions().dtype(torch::kFloat16).device(_in_feats.device()); + // at::Tensor _out_feats = + // torch::empty({num_in_feats, _kernel.size(0)}, options); + int num_out_feats = _out_feats.size(-2); + int num_out_channels = _out_feats.size(-1); + + + auto out_feats = reinterpret_cast(_out_feats.data_ptr()); + + if (num_out_feats > 128) + { + constexpr int CTA_M = 128; + constexpr int CTA_N = 128; + constexpr int CTA_K = 64; + constexpr int WARP_M = 128; + constexpr int WARP_N = 32; + constexpr int WARP_K = 64; + constexpr int STAGES = 3; + KERNEL_LAUNCH_CODE + } + else + { + constexpr int CTA_M = 64; + constexpr int CTA_N = 64; + constexpr int CTA_K = 64; + constexpr int WARP_M = 32; + constexpr int WARP_N = 32; + constexpr int WARP_K = 64; + constexpr int STAGES = 6; + KERNEL_LAUNCH_CODE + } + return ; +} \ No newline at end of file diff --git a/llm-awq/awq/kernels/csrc/w8a8/w8a8_gemm_cuda.h b/llm-awq/awq/kernels/csrc/w8a8/w8a8_gemm_cuda.h new file mode 100644 index 0000000000000000000000000000000000000000..6355fb99ec4d6fbd956df23c008e4236060512b8 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/w8a8/w8a8_gemm_cuda.h @@ -0,0 +1,4 @@ +#include + +void w8a8_gemm_forward_cuda(torch::Tensor _in_feats, torch::Tensor _kernel, torch::Tensor _wscales, torch::Tensor _ascales, torch::Tensor _out_feats); +void w8a8_gemm_fuse_bias_forward_cuda(torch::Tensor _in_feats, torch::Tensor _kernel, torch::Tensor _wscales, torch::Tensor _ascales, torch::Tensor _out_feats, torch::Tensor _bias); \ No newline at end of file diff --git a/llm-awq/awq/kernels/setup.py b/llm-awq/awq/kernels/setup.py new file mode 100644 index 0000000000000000000000000000000000000000..a44e1a5bb2d88352e6da5de690e8a2964ffbb9ac --- /dev/null +++ b/llm-awq/awq/kernels/setup.py @@ -0,0 +1,51 @@ +from setuptools import find_packages, setup +from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CppExtension + + +extra_compile_args = { + "cxx": ["-g", "-O3", "-fopenmp", "-lgomp", "-std=c++17", "-DENABLE_BF16"], + "nvcc": [ + "-O3", + "-std=c++17", + "-DENABLE_BF16", # TODO + "-U__CUDA_NO_HALF_OPERATORS__", + "-U__CUDA_NO_HALF_CONVERSIONS__", + "-U__CUDA_NO_BFLOAT16_OPERATORS__", + "-U__CUDA_NO_BFLOAT16_CONVERSIONS__", + "-U__CUDA_NO_BFLOAT162_OPERATORS__", + "-U__CUDA_NO_BFLOAT162_CONVERSIONS__", + "--expt-relaxed-constexpr", + "--expt-extended-lambda", + "--use_fast_math", + "--threads=8", + ], +} + +setup( + name="awq_inference_engine", + packages=find_packages(), + ext_modules=[ + CUDAExtension( + name="awq_inference_engine", + sources=[ + "csrc/pybind.cpp", + "csrc/quantization/gemm_cuda_gen.cu", + "csrc/quantization/gemv_cuda.cu", + "csrc/quantization_new/gemv/gemv_cuda.cu", + "csrc/quantization_new/gemm/gemm_cuda.cu", + "csrc/layernorm/layernorm.cu", + "csrc/position_embedding/pos_encoding_kernels.cu", + "csrc/attention/ft_attention.cpp", + "csrc/attention/decoder_masked_multihead_attention.cu", + "csrc/rope_new/fused_rope_with_pos.cu", + "csrc/w8a8/w8a8_gemm_cuda.cu", + "csrc/w8a8/quantization.cu", + "csrc/w8a8/act.cu", + "csrc/w8a8/layernorm.cu" + ], + extra_compile_args=extra_compile_args, + ), + ], + cmdclass={"build_ext": BuildExtension}, + install_requires=["torch"], +) diff --git a/llm-awq/awq/quantize/__init__.py b/llm-awq/awq/quantize/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..17ddad2a142490a81b6d44ebedd7629678d321c1 --- /dev/null +++ b/llm-awq/awq/quantize/__init__.py @@ -0,0 +1,2 @@ +from .w8a8_linear import * +from .smooth import * diff --git a/llm-awq/awq/quantize/__pycache__/__init__.cpython-311.pyc b/llm-awq/awq/quantize/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bcfd14b42f5778fb22665cb20db44206b3fa166b Binary files /dev/null and b/llm-awq/awq/quantize/__pycache__/__init__.cpython-311.pyc differ diff --git a/llm-awq/awq/quantize/__pycache__/w8a8_linear.cpython-311.pyc b/llm-awq/awq/quantize/__pycache__/w8a8_linear.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..15f83ba6d994799a071f95f953178bee9b0425d1 Binary files /dev/null and b/llm-awq/awq/quantize/__pycache__/w8a8_linear.cpython-311.pyc differ diff --git a/llm-awq/awq/quantize/auto_clip.py b/llm-awq/awq/quantize/auto_clip.py new file mode 100644 index 0000000000000000000000000000000000000000..0714bae50bfa37c8873df2a8ae9ae6a038bd4d04 --- /dev/null +++ b/llm-awq/awq/quantize/auto_clip.py @@ -0,0 +1,98 @@ +import torch +import torch.nn as nn +from .quantizer import pseudo_quantize_tensor +import gc + +__all__ = ["auto_clip_block"] + + +# weight quantization +@torch.no_grad() +def auto_clip_layer( + w, input_feat, n_bit, q_config, n_grid=20, max_shrink=0.5, n_sample_token=512 +): + assert w.dim() == 2 + org_w_shape = w.shape + # w [co, ci] -> [co, 1, n_group, group size] + # input_feat [n_token, ci] -> [1, n_token, n_group, group size] + group_size = ( + q_config["q_group_size"] if q_config["q_group_size"] > 0 else w.shape[1] + ) + input_feat = input_feat.view(-1, input_feat.shape[-1]) + input_feat = input_feat.reshape(1, input_feat.shape[0], -1, group_size) + input_feat = input_feat[:, 0 :: input_feat.shape[1] // n_sample_token] + w = w.reshape(w.shape[0], 1, -1, group_size) + + oc_batch_size = 256 if w.shape[0] % 256 == 0 else 64 # prevent OOM + assert w.shape[0] % oc_batch_size == 0 + w_all = w + best_max_val_all = [] + + for i_b in range(w.shape[0] // oc_batch_size): + w = w_all[i_b * oc_batch_size : (i_b + 1) * oc_batch_size] + + org_max_val = w.abs().amax(dim=-1, keepdim=True) # co, 1, n_group, 1 + + best_max_val = org_max_val.clone() + min_errs = torch.ones_like(org_max_val) * 1e9 + input_feat = input_feat.to(w.device) + org_out = (input_feat * w).sum(dim=-1) # co, n_token, n_group + + for i_s in range(int(max_shrink * n_grid)): + max_val = org_max_val * (1 - i_s / n_grid) + min_val = -max_val + cur_w = torch.clamp(w, min_val, max_val) + q_w = pseudo_quantize_tensor(cur_w, n_bit=n_bit, **q_config) + cur_out = (input_feat * q_w).sum(dim=-1) + + # co, 1, n_group, 1 + err = (cur_out - org_out).pow(2).mean(dim=1).view(min_errs.shape) + del cur_w + del cur_out + cur_best_idx = err < min_errs + min_errs[cur_best_idx] = err[cur_best_idx] + best_max_val[cur_best_idx] = max_val[cur_best_idx] + best_max_val_all.append(best_max_val) + + best_max_val = torch.cat(best_max_val_all, dim=0) + + del input_feat + del org_out + gc.collect() + torch.cuda.empty_cache() + return best_max_val.squeeze(1) + + +@torch.no_grad() +def auto_clip_block(module, w_bit, q_config, input_feat): + named_linears = { + name: m for name, m in module.named_modules() if isinstance(m, nn.Linear) + } + + clip_list = [] + for name in named_linears: + # due to qk bmm, it is hard to clip precisely + if any([_ in name for _ in ["q_", "k_", "query", "key", "Wqkv"]]): + continue + named_linears[name].cuda() + max_val = auto_clip_layer( + named_linears[name].weight, input_feat[name], n_bit=w_bit, q_config=q_config + ) + clip_list.append((name, max_val)) + named_linears[name].cpu() + return clip_list + + +@torch.no_grad() +def apply_clip(module, clip_list): + from ..utils.module import get_op_by_name + + for name, max_val in clip_list: + layer = get_op_by_name(module, name) + layer.cuda() + max_val = max_val.to(layer.weight.device).to(layer.weight.dtype) + org_shape = layer.weight.shape + layer.weight.data = layer.weight.data.reshape(*max_val.shape[:2], -1) + layer.weight.data = torch.clamp(layer.weight.data, -max_val, max_val) + layer.weight.data = layer.weight.data.reshape(org_shape) + layer.cpu() diff --git a/llm-awq/awq/quantize/auto_scale.py b/llm-awq/awq/quantize/auto_scale.py new file mode 100644 index 0000000000000000000000000000000000000000..fe604908ad9aa6a97b597d6dfa43c7b427913b17 --- /dev/null +++ b/llm-awq/awq/quantize/auto_scale.py @@ -0,0 +1,480 @@ +import gc +import torch +import torch.nn as nn + +from transformers.models.bloom.modeling_bloom import BloomBlock, BloomGelu +from transformers.models.opt.modeling_opt import OPTDecoderLayer +from transformers.models.llama.modeling_llama import LlamaDecoderLayer, LlamaRMSNorm +from transformers.activations import GELUActivation +from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm, Qwen2DecoderLayer + +from .qmodule import ScaledActivation +from ..utils.module import get_op_by_name, get_op_name, set_op_by_name + +__all__ = ["auto_scale_block", "apply_scale"] + + +@torch.no_grad() +def get_weight_scale(weight, q_group_size=-1): + org_shape = weight.shape + if q_group_size > 0: + weight = weight.view(-1, q_group_size) + scale = weight.abs() / weight.abs().amax(dim=1, keepdim=True) + scale = scale.view(org_shape) + scale = scale.mean(0) + return scale + + +@torch.no_grad() +def get_act_scale(x): + return x.abs().view(-1, x.shape[-1]).mean(0) + + +@torch.no_grad() +def scale_ln_fcs(ln, fcs, scales): + if not isinstance(fcs, list): + fcs = [fcs] + + scales = scales.to(ln.weight.device).to(ln.weight.dtype) + + ln.weight.div_(scales) + if hasattr(ln, "bias") and ln.bias is not None: + ln.bias.div_(scales) + + for fc in fcs: + fc.weight.mul_(scales.view(1, -1)) + + for p in ln.parameters(): + assert torch.isnan(p).sum() == 0 + for fc in fcs: + for p in fc.parameters(): + assert torch.isnan(p).sum() == 0 + + +@torch.no_grad() +def scale_fc_fc(fc1, fc2, scales): + assert isinstance(fc1, nn.Linear) + assert isinstance(fc2, nn.Linear) + # assert fc1.out_features == fc2.in_features + + scales = scales.to(fc1.weight.device).to(fc1.weight.dtype) + + # fc1.weight.div_(scales.view(-1, 1)) + fc1.weight[-scales.size(0) :].div_(scales.view(-1, 1)) + if fc1.bias is not None: + fc1.bias.div_(scales.view(-1)) + + fc2.weight.mul_(scales.view(1, -1)) + + for p in fc1.parameters(): + assert torch.isnan(p).sum() == 0 + for p in fc2.parameters(): + assert torch.isnan(p).sum() == 0 + + +@torch.no_grad() +def scale_gelu_fc(gelu, fc, scales): + assert isinstance(gelu, (nn.GELU, BloomGelu, GELUActivation)) + assert isinstance(fc, nn.Linear) + + fc.weight.mul_(scales.view(1, -1).to(fc.weight.device).to(fc.weight.dtype)) + + for p in fc.parameters(): + assert torch.isnan(p).sum() == 0 + + +@torch.no_grad() +def auto_scale_block(module, module_kwargs, w_bit, q_config, input_feat): + from .quantizer import pseudo_quantize_tensor + + # firstly, get the weight quantize function + if w_bit is not None: + + def w_quantize_func(p): + return pseudo_quantize_tensor( + p, + n_bit=w_bit, + **q_config, + ).detach() + + else: + + def w_quantize_func(p): + return p + + if "use_cache" in module_kwargs: + module_kwargs.pop("use_cache") + + # find the best scale ratio + def _search_module_scale(block, linears2scale: list, x, kwargs={}): + # w: co, ci + # x: n, ci + x = x.to(next(block.parameters()).device) + with torch.no_grad(): + org_out = block(x, **kwargs) + if isinstance(org_out, tuple): + org_out = org_out[0] + + x_max = get_act_scale(x) + + best_error = float("inf") + best_ratio = -1 + best_scales = None + + n_grid = 20 + history = [] + + org_sd = {k: v.cpu() for k, v in block.state_dict().items()} + for ratio in range(n_grid): + ratio = ratio * 1 / n_grid + scales = x_max.pow(ratio).clamp(min=1e-4).view(-1) + scales = scales / (scales.max() * scales.min()).sqrt() + for fc in linears2scale: + fc.weight.mul_(scales.view(1, -1).to(fc.weight.device)) + fc.weight.data = w_quantize_func(fc.weight.data) / (scales.view(1, -1)) + out = block(x, **kwargs) + if isinstance(out, tuple): + out = out[0] + + loss = ( + (org_out - out).float().pow(2).mean().item() + ) # float prevents overflow + history.append(loss) + is_best = loss < best_error + if is_best: + best_error = loss + best_ratio = ratio + best_scales = scales + block.load_state_dict(org_sd) + if best_ratio == -1: + print(history) + raise Exception + # print(best_ratio) + best_scales = best_scales.view(-1) + + assert torch.isnan(best_scales).sum() == 0, best_scales + return best_scales.detach() + + def _auto_get_scale(prev_op, layers, inp, module2inspect=None, kwargs={}): + # module2inspect: if given, we will check the output diff of this module instead of layers + if module2inspect is None: + assert len(layers) == 1 + module2inspect = layers[0] + + scales = _search_module_scale(module2inspect, layers, inp, kwargs) + scales = scales.detach().cpu() + # prev_op_name, [layer_name], scale + return ( + get_op_name(module, prev_op), + tuple([get_op_name(module, m) for m in layers]), + scales, + ) + + scales_list = [] # return the searched scales + + if isinstance(module, OPTDecoderLayer): + # attention input + scales_list.append( + _auto_get_scale( + prev_op=module.self_attn_layer_norm, + layers=[ + module.self_attn.q_proj, + module.self_attn.k_proj, + module.self_attn.v_proj, + ], + inp=input_feat["self_attn.q_proj"], + module2inspect=module.self_attn, + kwargs=module_kwargs, + ) + ) + # attn out + scales_list.append( + _auto_get_scale( + prev_op=module.self_attn.v_proj, + layers=[module.self_attn.out_proj], + inp=input_feat["self_attn.out_proj"], + ) + ) + # fc1 + scales_list.append( + _auto_get_scale( + prev_op=module.final_layer_norm, + layers=[module.fc1], + inp=input_feat["fc1"], + ) + ) + # fc2 + scales_list.append( + _auto_get_scale( + prev_op=module.fc1, + layers=[module.fc2], + inp=input_feat["fc2"], + ) + ) + + elif isinstance(module, (LlamaDecoderLayer, Qwen2DecoderLayer)): + # attention input + scales_list.append( + _auto_get_scale( + prev_op=module.input_layernorm, + layers=[ + module.self_attn.q_proj, + module.self_attn.k_proj, + module.self_attn.v_proj, + ], + inp=input_feat["self_attn.q_proj"], + module2inspect=module.self_attn, + kwargs=module_kwargs, + ) + ) + # attn out + # Please refer to https://github.com/mit-han-lab/llm-awq/pull/67#issue-1850622696 + if module.self_attn.v_proj.weight.shape == module.self_attn.o_proj.weight.shape: + scales_list.append( + _auto_get_scale( + prev_op=module.self_attn.v_proj, + layers=[module.self_attn.o_proj], + inp=input_feat["self_attn.o_proj"], + ) + ) + # fc1 + scales_list.append( + _auto_get_scale( + prev_op=module.post_attention_layernorm, + layers=[module.mlp.gate_proj, module.mlp.up_proj], + inp=input_feat["mlp.gate_proj"], + module2inspect=module.mlp, + ) + ) + # fc2 + scales_list.append( + _auto_get_scale( + prev_op=module.mlp.up_proj, + layers=[module.mlp.down_proj], + inp=input_feat["mlp.down_proj"], + ) + ) + + elif isinstance(module, BloomBlock): + # attention input + scales_list.append( + _auto_get_scale( + prev_op=module.input_layernorm, + layers=[module.self_attention.query_key_value], + inp=input_feat["self_attention.query_key_value"], + module2inspect=module, + kwargs=module_kwargs, + ) + ) + # attn out + # Please refer to https://github.com/mit-han-lab/llm-awq/issues/2#issuecomment-1606297469 + """ + scales_list.append(_auto_get_scale( + prev_op=module.self_attention.query_key_value, + layers=[module.self_attention.dense], + inp=input_feat['self_attention.dense'], + )) + """ + # fc1 + scales_list.append( + _auto_get_scale( + prev_op=module.post_attention_layernorm, + layers=[module.mlp.dense_h_to_4h], + inp=input_feat["mlp.dense_h_to_4h"], + module2inspect=module, + kwargs=module_kwargs, + ) + ) + # fc2 + scales_list.append( + _auto_get_scale( + prev_op=module.mlp.gelu_impl, + layers=[module.mlp.dense_4h_to_h], + inp=input_feat["mlp.dense_4h_to_h"], + ) + ) + elif "mpt" in str(module.__class__).lower(): + # attention input + scales_list.append( + _auto_get_scale( + prev_op=module.norm_1, + layers=[module.attn.Wqkv], + inp=input_feat["attn.Wqkv"], + module2inspect=module.attn, + kwargs=module_kwargs, + ) + ) + + # attn out + scales_list.append( + _auto_get_scale( + prev_op=module.attn.Wqkv, + layers=[module.attn.out_proj], + inp=input_feat["attn.out_proj"], + ) + ) + # fc1 + scales_list.append( + _auto_get_scale( + prev_op=module.norm_2, + layers=[module.ffn.up_proj], + inp=input_feat["ffn.up_proj"], + module2inspect=module.ffn, + ) + ) + # fc2 + scales_list.append( + _auto_get_scale( + prev_op=module.ffn.act, + layers=[module.ffn.down_proj], + inp=input_feat["ffn.down_proj"], + ) + ) + + elif "falcon" in str(module.__class__).lower(): + # attn out + # Haotian: TBD: need to handle repeated scales for MQ + """ + scales_list.append(_auto_get_scale( + prev_op=module.self_attention.query_key_value, + layers=[module.self_attention.dense], + inp=input_feat['self_attention.dense'], + )) + """ + # fc1, as long as it is scaled, everything is screwed up + if "falcon-7b" in str(module.__class__).lower(): + scales_list.append( + _auto_get_scale( + prev_op=module.input_layernorm, + layers=[ + module.mlp.dense_h_to_4h, + module.self_attention.query_key_value, + ], + inp=input_feat["self_attention.query_key_value"], + module2inspect=module, + kwargs=module_kwargs, + ) + ) + elif "falcon-40b" in str(module.__class__).lower(): + scales_list.append( + _auto_get_scale( + prev_op=module.ln_attn, + layers=[module.self_attention.query_key_value], + inp=input_feat["self_attention.query_key_value"], + module2inspect=module, + kwargs=module_kwargs, + ) + ) + scales_list.append( + _auto_get_scale( + prev_op=module.ln_mlp, + layers=[module.mlp.dense_h_to_4h], + inp=input_feat["mlp.dense_h_to_4h"], + module2inspect=module, + kwargs=module_kwargs, + ) + ) + else: + raise NotImplementedError( + "Unknown Falcon architecture, currently only falcon-7b and falcon-40b are supported" + ) + # fc2 + scales_list.append( + _auto_get_scale( + prev_op=module.mlp.act, + layers=[module.mlp.dense_4h_to_h], + inp=input_feat["mlp.dense_4h_to_h"], + ) + ) + elif "bigcode" in str(module.__class__).lower(): + scales_list.append( + _auto_get_scale( + prev_op=module.ln_1, + layers=[module.attn.c_attn], + inp=input_feat["attn.c_attn"], + module2inspect=module.attn, + kwargs=module_kwargs, + ) + ) + # fc1 + scales_list.append( + _auto_get_scale( + prev_op=module.ln_2, + layers=[module.mlp.c_fc], + inp=input_feat["mlp.c_fc"], + module2inspect=module.mlp, + ) + ) + # fc2 + scales_list.append( + _auto_get_scale( + prev_op=module.mlp.act, + layers=[module.mlp.c_proj], + inp=input_feat["mlp.c_proj"], + ) + ) + elif "neox" in str(module.__class__).lower(): + scales_list.append( + _auto_get_scale( + prev_op=module.input_layernorm, + layers=[module.attention.query_key_value], + inp=input_feat["attention.query_key_value"], + module2inspect=module.attention, + kwargs=module_kwargs, + ) + ) + # fc1 + scales_list.append( + _auto_get_scale( + prev_op=module.post_attention_layernorm, + layers=[module.mlp.dense_h_to_4h], + inp=input_feat["mlp.dense_h_to_4h"], + module2inspect=module.mlp, + ) + ) + # fc2 + scales_list.append( + _auto_get_scale( + prev_op=module.mlp.act, + layers=[module.mlp.dense_4h_to_h], + inp=input_feat["mlp.dense_4h_to_h"], + ) + ) + else: + raise NotImplementedError(f"{type(module)} not supported yet!") + + return scales_list + + +def apply_scale(module, scales_list, input_feat_dict=None): + for prev_op_name, layer_names, scales in scales_list: + prev_op = get_op_by_name(module, prev_op_name) + layers = [get_op_by_name(module, name) for name in layer_names] + + prev_op.cuda() + for layer in layers: + layer.cuda() + scales.cuda() + + if isinstance(prev_op, nn.Linear): + assert len(layers) == 1 + scale_fc_fc(prev_op, layers[0], scales) + elif isinstance(prev_op, (nn.LayerNorm, LlamaRMSNorm, Qwen2RMSNorm)): + scale_ln_fcs(prev_op, layers, scales) + elif isinstance(prev_op, (nn.GELU, BloomGelu, GELUActivation, nn.SiLU)): + new_module = ScaledActivation(prev_op, scales) + set_op_by_name(module, prev_op_name, new_module) + scale_gelu_fc(prev_op, layers[0], scales) + else: + raise NotImplementedError(f"prev_op {type(prev_op)} not supported yet!") + + # apply the scaling to input feat if given; prepare it for clipping + if input_feat_dict is not None: + for layer_name in layer_names: + inp = input_feat_dict[layer_name] + inp.div_(scales.view(1, -1).to(inp.device).to(inp.dtype)) + + prev_op.cpu() + for layer in layers: + layer.cpu() + scales.cpu() diff --git a/llm-awq/awq/quantize/pre_quant.py b/llm-awq/awq/quantize/pre_quant.py new file mode 100644 index 0000000000000000000000000000000000000000..23a5ecbe11befb977a79535701802703be31042e --- /dev/null +++ b/llm-awq/awq/quantize/pre_quant.py @@ -0,0 +1,254 @@ +import torch +import torch.nn as nn +import tqdm +import gc +import functools +from collections import defaultdict +from typing import List + +from transformers.models.bloom.modeling_bloom import BloomForCausalLM +from transformers.models.opt.modeling_opt import OPTForCausalLM +from transformers.models.llama.modeling_llama import LlamaForCausalLM +try: + from tinychat.models import LlavaLlamaForCausalLM +except ImportError as e: + pass + +from transformers.models.qwen2.modeling_qwen2 import Qwen2ForCausalLM + +from .auto_scale import auto_scale_block, apply_scale +from .auto_clip import auto_clip_block, apply_clip + +__all__ = ["run_awq"] + + +def get_named_linears(module): + return {name: m for name, m in module.named_modules() if isinstance(m, nn.Linear)} + + +def get_blocks(model): + if model.__class__.__name__ in ("LlamaForCausalLM", "Qwen2ForCausalLM"): + layers = model.model.layers + elif model.__class__.__name__ == "InternVL3": + layers = model.language_model.model.layers + # layers = [model.language_model.model.layers, model.vision_model.encoder.layers] + elif model.__class__.__name__ == "LlavaLlamaForCausalLM": + # layers = [model.model.layers, model.model.vision_tower.vision_tower.vision_model.encoder.layers] + layers = model.model.layers + elif isinstance(model, OPTForCausalLM): + layers = model.model.decoder.layers + elif isinstance(model, BloomForCausalLM): + layers = model.transformer.h + elif "mpt" in str(model.__class__).lower(): + layers = model.transformer.blocks + elif "falcon" in str(model.__class__).lower(): + layers = model.transformer.h + elif "bigcode" in str(model.__class__).lower(): + layers = model.transformer.h + elif "neox" in str(model.__class__).lower(): + layers = model.gpt_neox.layers + elif model.__class__.__name__ == "LlavaLlamaModel": + layers = model.llm.model.layers + else: + raise NotImplementedError(type(model)) + return layers + + +def move_embed(model, device): + if isinstance(model, (LlamaForCausalLM, Qwen2ForCausalLM)): + model.model.embed_tokens = model.model.embed_tokens.to(device) + model.model.rotary_emb = model.model.rotary_emb.to(device) + elif model.__class__.__name__ == "InternVL3": + model.language_model.model.embed_tokens = ( + model.language_model.model.embed_tokens.to(device) + ) + model.language_model.model.rotary_emb = ( + model.language_model.model.rotary_emb.to(device) + ) + model.vision_model.embeddings.to(device) + elif isinstance(model, LlavaLlamaForCausalLM): + model.model.embed_tokens = model.model.embed_tokens.to(device) + model.model.vision_tower.vision_tower.vision_model.embeddings.to(device) + elif isinstance(model, OPTForCausalLM): + model.model.decoder.embed_tokens = model.model.decoder.embed_tokens.to(device) + model.model.decoder.embed_positions = model.model.decoder.embed_positions.to( + device + ) + elif isinstance(model, BloomForCausalLM): + model.transformer.word_embeddings = model.transformer.word_embeddings.to(device) + model.transformer.word_embeddings_layernorm = ( + model.transformer.word_embeddings_layernorm.to(device) + ) + elif "mpt" in str(model.__class__).lower(): + model.transformer.wte = model.transformer.wte.to(device) + model.transformer.emb_drop = model.transformer.emb_drop.to(device) + elif "falcon" in str(model.__class__).lower(): + model.transformer.word_embeddings = model.transformer.word_embeddings.to(device) + elif "bigcode" in str(model.__class__).lower(): + model.transformer.wte = model.transformer.wte.to(device) + model.transformer.wpe = model.transformer.wpe.to(device) + model.transformer.drop = model.transformer.drop.to(device) + elif "neox" in str(model.__class__).lower(): + model.gpt_neox.embed_in = model.gpt_neox.embed_in.to(device) + model.gpt_neox.emb_dropout = model.gpt_neox.emb_dropout.to(device) + model.embed_out = model.embed_out.to(device) + elif "llavallamamodel" in str(model.__class__).lower(): + model.llm.model.embed_tokens = model.llm.model.embed_tokens.to(device) + else: + raise NotImplementedError(type(model)) + + +@torch.no_grad() +def run_awq( + model, + enc, + w_bit, + q_config, + n_samples=512, + seqlen=512, + auto_scale=True, + mse_range=True, + # some configs for ablation study + calib_data="pileval", +): + from ..utils.calib_data import get_calib_dataset + from ..utils.module import append_str_prefix, get_op_name + + if "bigcode" in str(model.__class__).lower(): + # otherwise attention_mask will always be on cpu. + model.transformer.bias = model.transformer.bias.to("cuda") + + layers = get_blocks(model) + + samples = get_calib_dataset( + data=calib_data, tokenizer=enc, n_samples=n_samples, block_size=seqlen + ) + samples = torch.cat(samples, dim=0) + + inps = [] + layer_kwargs = {} + + layers[0] = layers[0].cuda() + move_embed(model, "cuda") + + # get input and kwargs to layer 0 + # with_kwargs is only supported in PyTorch 2.0 + # use this Catcher hack for now + class Catcher(nn.Module): + def __init__(self, module): + super().__init__() + self.module = module + + def forward(self, inp, **kwargs): + inps.append(inp) + layer_kwargs.update(kwargs) + raise ValueError # early exit to break later inference + + # patch layer 0 to catch input and kwargs + layers[0] = Catcher(layers[0]) + try: + if model.__class__.__name__ == "LlavaLlamaModel": + model.llm(samples.to(next(model.parameters()).device)) + elif model.__class__.__name__ == "InternVL3": + model.language_model(samples.to(next(model.parameters()).device)) + else: + model(samples.to(next(model.parameters()).device)) + except ValueError: # work with early exit + pass + del samples + layers[0] = layers[0].module # restore + inps = inps[0] + + layers[0] = layers[0].cpu() + move_embed(model, "cpu") + + gc.collect() + torch.cuda.empty_cache() + + awq_results = { + "scale": [], + "clip": [], + } + + # solve layer by layer + for i in tqdm.tqdm(range(len(layers)), desc="Running AWQ..."): + layer = layers[i] + layer = layer.cuda() + named_linears = get_named_linears(layer) + + # firstly, get input features of all linear layers + def cache_input_hook(m, x, y, name, feat_dict): + x = x[0] + x = x.detach().cpu() + feat_dict[name].append(x) + + input_feat = defaultdict(list) + handles = [] + for name in named_linears: + handles.append( + named_linears[name].register_forward_hook( + functools.partial(cache_input_hook, name=name, feat_dict=input_feat) + ) + ) + inps = inps.to(next(layer.parameters()).device) # in case multi-gpu + # get output as next layer's input + inps = layer(inps, **layer_kwargs)[0] + for h in handles: + h.remove() + # now solve for scaling and clipping + input_feat = {k: torch.cat(v, dim=0) for k, v in input_feat.items()} + + # Clear GPU memory + torch.cuda.empty_cache() + + if ( + auto_scale + ): # if it applies, we should also modify the input_feat with scales + scales_list = auto_scale_block( + layer, + layer_kwargs, + w_bit=w_bit, + q_config=q_config, + input_feat=input_feat, + ) + # apply_scale(layer, scales_list, input_feat_dict=input_feat) + apply_scale(layers[i], scales_list, input_feat_dict=input_feat) + # append prefix to make names global + awq_results["scale"] += append_str_prefix( + scales_list, get_op_name(model, layer) + "." + ) + + # Clear GPU memory + torch.cuda.empty_cache() + # for line in torch.cuda.memory_summary().splitlines(): + # if "Allocated" in line: + # print(line) + + if mse_range: + clip_list = auto_clip_block( + layer, + w_bit=w_bit, + q_config=q_config, + input_feat=input_feat, + ) + apply_clip(layer, clip_list) + # append prefix to make names global + awq_results["clip"] += append_str_prefix( + clip_list, get_op_name(model, layer) + "." + ) + + layer = layer.cpu() + # Haotian: check activation replacement + del input_feat + gc.collect() + torch.cuda.empty_cache() + # for line in torch.cuda.memory_summary().splitlines(): + # if "Allocated" in line: + # print(line) + + return awq_results + + +def apply_awq(model, awq_results): + apply_scale(model, awq_results["scale"]) + apply_clip(model, awq_results["clip"]) diff --git a/llm-awq/awq/quantize/quantizer.py b/llm-awq/awq/quantize/quantizer.py new file mode 100644 index 0000000000000000000000000000000000000000..7d15191651545f504ace5e0f284525421e423bf1 --- /dev/null +++ b/llm-awq/awq/quantize/quantizer.py @@ -0,0 +1,165 @@ +import torch +import torch.nn as nn +from tqdm import tqdm +import gc +from .qmodule import ScaledActivation +from ..utils.module import set_op_by_name + +from transformers.models.bloom.modeling_bloom import BloomBlock + +EMBEDDING_KEYWORDS = ["embed"] +LM_HEAD_KEYWORDS = ["lm_head", "embed_out", "output"] + + +def scale_activations(module): + param = next(module.parameters()) + dtype = param.dtype + device = param.device + if isinstance(module, BloomBlock): + if isinstance(module.mlp.gelu_impl, ScaledActivation): + return + c = module.mlp.dense_h_to_4h.out_features + act = ScaledActivation( + module.mlp.gelu_impl, torch.ones(c, dtype=dtype, device=device) + ) + set_op_by_name(module, "mlp.gelu_impl", act) + elif "mptblock" in str(module.__class__.__name__).lower(): + if isinstance(module.ffn.act, ScaledActivation): + return + c = module.ffn.up_proj.out_features + act = ScaledActivation( + module.ffn.act, torch.ones(c, dtype=dtype, device=device) + ) + set_op_by_name(module, "ffn.act", act) + elif "falcon" in str(module.__class__).lower(): + if isinstance(module.mlp.act, ScaledActivation): + return + c = module.mlp.dense_h_to_4h.out_features + act = ScaledActivation( + module.mlp.act, torch.ones(c, dtype=dtype, device=device) + ) + set_op_by_name(module, "mlp.act", act) + elif "bigcode" in str(module.__class__).lower(): + if isinstance(module.mlp.act, ScaledActivation): + return + c = module.mlp.c_proj.out_features + act = ScaledActivation( + module.mlp.act, torch.ones(c, dtype=dtype, device=device) + ) + set_op_by_name(module, "mlp.act", act) + elif "neox" in str(module.__class__).lower(): + if isinstance(module.mlp.act, ScaledActivation): + return + c = module.mlp.dense_h_to_4h.out_features + act = ScaledActivation( + module.mlp.act, torch.ones(c, dtype=dtype, device=device) + ) + set_op_by_name(module, "mlp.act", act) + + +# core quantization method (simulated quantization) +def pseudo_quantize_tensor( + w, n_bit=8, zero_point=True, q_group_size=-1, inplace=False, get_scale_zp=False +): + org_w_shape = w.shape + if q_group_size > 0: + assert org_w_shape[-1] % q_group_size == 0 + w = w.reshape(-1, q_group_size) + assert w.dim() == 2 + if zero_point: + max_val = w.amax(dim=1, keepdim=True) + min_val = w.amin(dim=1, keepdim=True) + max_int = 2**n_bit - 1 + min_int = 0 + scales = (max_val - min_val).clamp(min=1e-5) / max_int + zeros = (-torch.round(min_val / scales)).clamp_(min_int, max_int) + else: # we actually never used this + assert min_val is None + max_val = w.abs().amax(dim=1, keepdim=True) + max_val = max_val.clamp(min=1e-5) + max_int = 2 ** (n_bit - 1) - 1 + min_int = -(2 ** (n_bit - 1)) + scales = max_val / max_int + zeros = 0 + + assert torch.isnan(scales).sum() == 0 + assert torch.isnan(w).sum() == 0 + + if inplace: + ( + (w.div_(scales).round_().add_(zeros)).clamp_(min_int, max_int).sub_(zeros) + ).mul_(scales) + else: + w = ( + torch.clamp(torch.round(w / scales) + zeros, min_int, max_int) - zeros + ) * scales + assert torch.isnan(w).sum() == 0 + + w = w.reshape(org_w_shape) + + if get_scale_zp: + return w, scales.view(w.shape[0], -1), zeros.view(w.shape[0], -1) + else: + return w + + +@torch.no_grad() +def pseudo_quantize_model_weight( + model, + w_bit, + q_config, +): + from .pre_quant import get_blocks, get_named_linears + + layers = get_blocks(model) + for i in tqdm(range(len(layers)), desc="pseudo weight quantization..."): + named_linears = get_named_linears(layers[i]) + for n, m in named_linears.items(): + m.cuda() + m.weight.data = pseudo_quantize_tensor( + m.weight.data, n_bit=w_bit, **q_config + ) + m.cpu() + + +@torch.no_grad() +def real_quantize_model_weight(model, w_bit, q_config, init_only=False): + from .qmodule import WQLinear + from .pre_quant import get_blocks, get_named_linears + + assert q_config["zero_point"], "We only support zero_point quantization now." + + layers = get_blocks(model) + for i in tqdm( + range(len(layers)), + desc="real weight quantization..." + ("(init only)" if init_only else ""), + ): + layer = layers[i] + named_linears = get_named_linears(layer) + scale_activations(layer) + + for name, module in named_linears.items(): + if init_only: + q_linear = WQLinear.from_linear( + module, w_bit, q_config["q_group_size"], True + ) + q_linear.to(next(layer.parameters()).device) + set_op_by_name(layer, name, q_linear) + else: + module.cuda() + module.weight.data, scales, zeros = pseudo_quantize_tensor( + module.weight.data, n_bit=w_bit, get_scale_zp=True, **q_config + ) + # scales = scales.t().contiguous() + # zeros = zeros.t().contiguous() + q_linear = WQLinear.from_linear( + module, w_bit, q_config["q_group_size"], False, scales, zeros + ) + module.cpu() + q_linear.to(next(layer.parameters()).device) + set_op_by_name(layer, name, q_linear) + torch.cuda.empty_cache() + gc.collect() + + torch.cuda.empty_cache() + gc.collect() diff --git a/llm-awq/awq/quantize/smooth.py b/llm-awq/awq/quantize/smooth.py new file mode 100644 index 0000000000000000000000000000000000000000..f4393f81edf2c0dbccb055c570c713ce3ebb3502 --- /dev/null +++ b/llm-awq/awq/quantize/smooth.py @@ -0,0 +1,246 @@ +# Adapted from SmoothQuant (https://github.com/mit-han-lab/smoothquant) and modified by Yuming Lou + + +import torch.nn as nn +try: + import llava + from llava.media import Image, Video + from llava.utils.media import extract_media + from llava.constants import DEFAULT_IMAGE_TOKEN + from llava.mm_utils import process_image, process_images +except ImportError: + print("VILA is not installed. Multimodal features will not be available. To activate, please install VILA at https://github.com/NVlabs/VILA.") + +import torch +from collections import defaultdict +from functools import partial +from tqdm import tqdm +import numpy as np +import functools + + +@torch.no_grad() +def get_act_scales(model, data): + num_samples = data.shape[0] + model.eval() + act_scales = {} + + def stat_tensor(name, tensor): + hidden_dim = tensor.shape[-1] + tensor = tensor.view(-1, hidden_dim).abs().detach() + comming_max = torch.max(tensor, dim=0)[0].float().cpu() + if name in act_scales: + act_scales[name] = torch.max(act_scales[name], comming_max) + else: + act_scales[name] = comming_max + + def stat_input_hook(m, x, y, name): + if isinstance(x, tuple): + x = x[0] + stat_tensor(name, x) + + hooks = [] + for name, m in model.named_modules(): + if isinstance(m, nn.Linear): + hooks.append( + m.register_forward_hook(functools.partial(stat_input_hook, name=name)) + ) + + for i in tqdm(range(num_samples)): + input = data[i : i + 1] + model(input) + + for h in hooks: + h.remove() + + return act_scales + + +@torch.no_grad() +def get_static_decoder_layer_scales( + model, + data, +): + num_samples = data.shape[1] + model.eval() + device = next(model.parameters()).device + + act_dict = defaultdict(dict) + + def stat_io_hook(m, x, y, name): + if isinstance(x, tuple): + x = x[0] + if name not in act_dict or "input" not in act_dict[name]: + act_dict[name]["input"] = x.detach().abs().max().item() + else: + act_dict[name]["input"] = max( + act_dict[name]["input"], x.detach().abs().max().item() + ) + if isinstance(y, tuple): + y = y[0] + if name not in act_dict or "output" not in act_dict[name]: + act_dict[name]["output"] = y.detach().abs().max().item() + else: + act_dict[name]["output"] = max( + act_dict[name]["output"], y.detach().abs().max().item() + ) + + hooks = [] + for name, m in model.named_modules(): + if isinstance(m, torch.nn.Linear): + hooks.append(m.register_forward_hook(partial(stat_io_hook, name=name))) + pbar = tqdm(range(num_samples)) + for i in pbar: + model(data[i : i + 1]) + mean_scale = np.mean([v["input"] for v in act_dict.values()]) + pbar.set_description(f"Mean input scale: {mean_scale:.2f}") + for hook in hooks: + hook.remove() + decoder_layer_scales = [] + for idx in range(model.config.num_hidden_layers): + scale_dict = {} + scale_dict["attn_input_scale"] = ( + act_dict[ + f"vision_tower.vision_model.encoder.layers.{idx}.self_attn.q_proj" + ]["input"] + / 127 + ) + scale_dict["q_output_scale"] = ( + act_dict[ + f"vision_tower.vision_model.encoder.layers.{idx}.self_attn.q_proj" + ]["output"] + / 127 + ) + scale_dict["k_output_scale"] = ( + act_dict[ + f"vision_tower.vision_model.encoder.layers.{idx}.self_attn.k_proj" + ]["output"] + / 127 + ) + scale_dict["v_output_scale"] = ( + act_dict[ + f"vision_tower.vision_model.encoder.layers.{idx}.self_attn.v_proj" + ]["output"] + / 127 + ) + scale_dict["out_input_scale"] = ( + act_dict[ + f"vision_tower.vision_model.encoder.layers.{idx}.self_attn.out_proj" + ]["input"] + / 127 + ) + scale_dict["fc1_input_scale"] = ( + act_dict[f"vision_tower.vision_model.encoder.layers.{idx}.mlp.fc1"]["input"] + / 127 + ) + scale_dict["fc2_input_scale"] = ( + act_dict[f"vision_tower.vision_model.encoder.layers.{idx}.mlp.fc2"]["input"] + / 127 + ) + decoder_layer_scales.append(scale_dict) + + return decoder_layer_scales, act_dict + + +def get_smooth_scale(model_path, media): + # Load model + model = llava.load(model_path, devices=[0]) + del model.llm + del model.mm_projector + torch.cuda.empty_cache() + model = model.cuda().eval() + prompt = [] + if media is not None: + for m in media or []: + if any(m.endswith(ext) for ext in [".jpg", ".jpeg", ".png"]): + m = Image(m) + elif any(m.endswith(ext) for ext in [".mp4", ".mkv", ".webm"]): + m = Video(m) + else: + raise ValueError(f"Unsupported media type: {m}") + prompt.append(m) + conversation = [{"from": "human", "value": prompt}] + media = extract_media(conversation, model.config) + for name in media: + if name == "image": + if ( + len(media["image"]) == 1 + and model.config.image_aspect_ratio == "dynamic" + ): + model.config.image_processor = model.vision_tower.image_processor + images = process_image( + media["image"][0], model.config, None, enable_dynamic_res=True + ).half() + conversation[0]["value"] = conversation[0]["value"].replace( + DEFAULT_IMAGE_TOKEN, f"{DEFAULT_IMAGE_TOKEN}\n" * images.shape[0] + ) + else: + images = process_images( + media["image"], model.vision_tower.image_processor, model.config + ).half() + media[name] = [image for image in images] + elif name == "video": + media[name] = [ + process_images( + images, model.vision_tower.image_processor, model.config + ).half() + for images in media[name] + ] + else: + raise ValueError(f"Unsupported media type: {name}") + images = torch.cat(media["video"], dim=1) + model.vision_tower = model.vision_tower.eval() + decoder_layer_scales = get_act_scales(model.vision_tower, images) + return decoder_layer_scales + + +@torch.no_grad() +def smooth_ln_fcs(ln, fcs, act_scales, alpha=0.5): + if not isinstance(fcs, list): + fcs = [fcs] + assert isinstance(ln, nn.LayerNorm) + for fc in fcs: + assert isinstance(fc, nn.Linear) + assert ln.weight.numel() == fc.in_features == act_scales.numel() + + device, dtype = fcs[0].weight.device, fcs[0].weight.dtype + act_scales = act_scales.to(device=device, dtype=dtype) + weight_scales = torch.cat( + [fc.weight.abs().max(dim=0, keepdim=True)[0] for fc in fcs], dim=0 + ) + weight_scales = weight_scales.max(dim=0)[0].clamp(min=1e-5) + + scales = ( + (act_scales.pow(alpha) / weight_scales.pow(1 - alpha)) + .clamp(min=1e-5) + .to(device) + .to(dtype) + ) + + ln.weight.div_(scales) + ln.bias.div_(scales) + + for fc in fcs: + fc.weight.mul_(scales.view(1, -1)) + + +@torch.no_grad() +def smooth_lm(model, scales, alpha=0.5): + if "siglip" in str(model.__class__).lower(): + num = 0 + for name, module in model.named_modules(): + if "siglipencoderlayer" in str(module.__class__).lower(): + attn_ln = module.layer_norm1 + qkv = [ + module.self_attn.q_proj, + module.self_attn.k_proj, + module.self_attn.v_proj, + ] + qkv_input_scales = scales[name + ".self_attn.q_proj"] + smooth_ln_fcs(attn_ln, qkv, qkv_input_scales, alpha) + + ffn_ln = module.layer_norm2 + fc1 = module.mlp.fc1 + fc1_input_scales = scales[name + ".mlp.fc1"] + smooth_ln_fcs(ffn_ln, fc1, fc1_input_scales, alpha) + num += 1 diff --git a/llm-awq/awq/utils/__init__.py b/llm-awq/awq/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/llm-awq/awq/utils/__pycache__/__init__.cpython-311.pyc b/llm-awq/awq/utils/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..466bfee723919b6b33537604eca3898fadaf74b6 Binary files /dev/null and b/llm-awq/awq/utils/__pycache__/__init__.cpython-311.pyc differ diff --git a/llm-awq/awq/utils/__pycache__/parallel.cpython-311.pyc b/llm-awq/awq/utils/__pycache__/parallel.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7616a7a63fda4388c326d62c7f6ee234c7712c69 Binary files /dev/null and b/llm-awq/awq/utils/__pycache__/parallel.cpython-311.pyc differ diff --git a/llm-awq/awq/utils/calib_data.py b/llm-awq/awq/utils/calib_data.py new file mode 100644 index 0000000000000000000000000000000000000000..96514ef12ee275ec8fca5ff5828a3707c9debb3a --- /dev/null +++ b/llm-awq/awq/utils/calib_data.py @@ -0,0 +1,32 @@ +import torch +from datasets import load_dataset + + +def get_calib_dataset(data="pileval", tokenizer=None, n_samples=512, block_size=512): + if data == "pileval": + dataset = load_dataset("mit-han-lab/pile-val-backup", split="validation") + else: + raise NotImplementedError + dataset = dataset.shuffle(seed=42) + samples = [] + n_run = 0 + for data in dataset: + line = data["text"] + line = line.strip() + line_encoded = tokenizer.encode(line) + if len(line_encoded) > 512: + continue + sample = torch.tensor([line_encoded]) + if sample.numel() == 0: + continue + samples.append(sample) + n_run += 1 + if n_run == n_samples: + break + # now concatenate all samples and split according to block size + cat_samples = torch.cat(samples, dim=1) + n_split = cat_samples.shape[1] // block_size + print(f" * Split into {n_split} blocks") + return [ + cat_samples[:, i * block_size : (i + 1) * block_size] for i in range(n_split) + ] diff --git a/llm-awq/awq/utils/module.py b/llm-awq/awq/utils/module.py new file mode 100644 index 0000000000000000000000000000000000000000..3c150b59e59927ad4632947307459acdd76a9437 --- /dev/null +++ b/llm-awq/awq/utils/module.py @@ -0,0 +1,39 @@ +def get_op_by_name(module, op_name): + # get the op by its name relative to the module + for name, m in module.named_modules(): + if name == op_name: + return m + raise ValueError(f"Cannot find op {op_name} in module {module}") + + +def set_op_by_name(layer, name, new_module): + levels = name.split(".") + if len(levels) > 1: + mod_ = layer + for l_idx in range(len(levels) - 1): + if levels[l_idx].isdigit(): + mod_ = mod_[int(levels[l_idx])] + else: + mod_ = getattr(mod_, levels[l_idx]) + setattr(mod_, levels[-1], new_module) + else: + setattr(layer, name, new_module) + + +def get_op_name(module, op): + # get the name of the op relative to the module + for name, m in module.named_modules(): + if m is op: + return name + raise ValueError(f"Cannot find op {op} in module {module}") + + +def append_str_prefix(x, prefix): + if isinstance(x, str): + return prefix + x + elif isinstance(x, tuple): + return tuple([append_str_prefix(y, prefix) for y in x]) + elif isinstance(x, list): + return [append_str_prefix(y, prefix) for y in x] + else: + return x diff --git a/llm-awq/awq/utils/parallel.py b/llm-awq/awq/utils/parallel.py new file mode 100644 index 0000000000000000000000000000000000000000..9aed0f657eae43e8b33c545f628b5d0da47cb02f --- /dev/null +++ b/llm-awq/awq/utils/parallel.py @@ -0,0 +1,28 @@ +import os +import torch +import gc + + +def auto_parallel(args): + model_size = args.model_path.split("-")[-1] + if model_size.endswith("m"): + model_gb = 1 + else: + model_gb = float(model_size[:-1]) + if model_gb < 20: + n_gpu = 1 + elif model_gb < 50: + n_gpu = 4 + else: + n_gpu = 8 + args.parallel = n_gpu > 1 + cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES", None) + if isinstance(cuda_visible_devices, str): + cuda_visible_devices = cuda_visible_devices.split(",") + else: + cuda_visible_devices = list(range(8)) + os.environ["CUDA_VISIBLE_DEVICES"] = ",".join( + [str(dev) for dev in cuda_visible_devices[:n_gpu]] + ) + print("CUDA_VISIBLE_DEVICES: ", os.environ["CUDA_VISIBLE_DEVICES"]) + return cuda_visible_devices diff --git a/llm-awq/examples/README.md b/llm-awq/examples/README.md new file mode 100644 index 0000000000000000000000000000000000000000..1d3a047ae5a53b4b4009d4fbf7837506608b19e6 --- /dev/null +++ b/llm-awq/examples/README.md @@ -0,0 +1,10 @@ +# AWQ Examples + +Here we provide two AWQ examples, applying to: +- [Vicuna-7B](https://github.com/lm-sys/FastChat), a chatbot with instruction-tuning +- [LLaVA-13B](https://github.com/lm-sys/FastChat), a visual LM for multi-modal applications like visual reasoning. +- [A simple conversion script](https://github.com/mit-han-lab/llm-awq/tree/main/examples/convert_to_hf.py) to convert llm-awq weights into HF format. + +Here are some example output from the two demos. You should able to observe memory saving when running the demos in 4-bit. Please check the notebooks for details. + +![overview](../figures/example_vis.jpg) diff --git a/llm-awq/examples/chat_demo.ipynb b/llm-awq/examples/chat_demo.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..c4bb3f1856b1aca3274fb6e3fe8474aa38ad6cf6 --- /dev/null +++ b/llm-awq/examples/chat_demo.ipynb @@ -0,0 +1,272 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# AWQ on Vicuna" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "In this notebook, we use Vicuna model to demonstrate the performance of AWQ on instruction-tuned models. We implement AWQ real-INT4 inference kernels, which are wrapped as Pytorch modules and can be easily used by existing models. We also provide a simple example to show how to use AWQ to quantize a model and save/load the quantized model checkpoint." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "In order to run this notebook, you need to install the following packages:\n", + "- [AWQ](https://github.com/mit-han-lab/llm-awq)\n", + "- [Pytorch](https://pytorch.org/)\n", + "- [Accelerate](https://github.com/huggingface/accelerate)\n", + "- [Transformers](https://github.com/huggingface/transformers)" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [], + "source": [ + "import torch\n", + "from accelerate import init_empty_weights, load_checkpoint_and_dispatch\n", + "from awq.quantize.quantizer import real_quantize_model_weight\n", + "from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig\n", + "from tinychat.demo import gen_params, stream_output\n", + "from tinychat.stream_generators import StreamGenerator\n", + "from tinychat.modules import make_quant_norm, make_quant_attn, make_fused_mlp\n", + "from tinychat.utils.prompt_templates import get_prompter\n", + "import os\n", + "# This demo only support single GPU for now\n", + "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Please get the Vicuna model from [FastChat](https://github.com/lm-sys/FastChat) and run the following command to generate a quantized model checkpoint first.\n", + "\n", + "```bash\n", + "mkdir quant_cache\n", + "python -m awq.entry --model_path [vicuna-7b_model_path] \\\n", + " --w_bit 4 --q_group_size 128 \\\n", + " --load_awq awq_cache/vicuna-7b-w4-g128.pt \\\n", + " --q_backend real --dump_quant quant_cache/vicuna-7b-w4-g128-awq.pt\n", + "```" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [], + "source": [ + "# model_path = \"\" # the path of vicuna-7b model\n", + "# load_quant_path = \"quant_cache/vicuna-7b-w4-g128-awq.pt\"\n", + "model_path = \"/data/llm/checkpoints/vicuna-hf/vicuna-7b\"\n", + "load_quant_path = \"/data/llm/checkpoints/vicuna-hf/vicuna-7b-awq-w4g128.pt\"" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We first load a empty model and replace all the linear layers with WQLinear layers. Then we load the quantized weights from the checkpoint. " + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [ + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "8b79a82b73ab4d9191ba54f5d0f8cb86", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "Loading checkpoint shards: 0%| | 0/2 [00:00=61.0"] +build-backend = "setuptools.build_meta" + +[project] +name = "awq" +version = "0.1.0" +description = "An efficient and accurate low-bit weight quantization(INT3/4) method for LLMs." +readme = "README.md" +requires-python = ">=3.8" +classifiers = [ + "Programming Language :: Python :: 3", + "License :: OSI Approved :: Apache Software License", +] +dependencies = [ + "accelerate==0.34.2", "sentencepiece", "tokenizers>=0.12.1", + "torch==2.3.0", "torchvision==0.18.0", + "transformers==4.46.0", + "lm_eval==0.3.0", "texttable", + "toml", "attributedict", + "protobuf", + "gradio==3.35.2", "gradio_client==0.2.9", + "fastapi", "uvicorn", + "pydantic==1.10.19" +] + +[tool.setuptools.packages.find] +exclude = ["results*", "scripts*", "examples*"] + +[tool.wheel] +exclude = ["results*", "scripts*", "examples*"] diff --git a/lm-evaluation-harness/lm_eval/models/__pycache__/anthropic_llms.cpython-311.pyc b/lm-evaluation-harness/lm_eval/models/__pycache__/anthropic_llms.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e7321ae87797a94760e6cd533fc4ca34822adbcb Binary files /dev/null and b/lm-evaluation-harness/lm_eval/models/__pycache__/anthropic_llms.cpython-311.pyc differ diff --git a/lm-evaluation-harness/lm_eval/models/__pycache__/hf_vlms.cpython-311.pyc b/lm-evaluation-harness/lm_eval/models/__pycache__/hf_vlms.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3645017e8ecd6cccbaef1efb536b4a4989c6b2ba Binary files /dev/null and b/lm-evaluation-harness/lm_eval/models/__pycache__/hf_vlms.cpython-311.pyc differ diff --git a/lm-evaluation-harness/lm_eval/models/__pycache__/mamba_lm.cpython-311.pyc b/lm-evaluation-harness/lm_eval/models/__pycache__/mamba_lm.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f857b833fc4bddbd662e61f96eff2dc5d927c91f Binary files /dev/null and b/lm-evaluation-harness/lm_eval/models/__pycache__/mamba_lm.cpython-311.pyc differ diff --git a/lm-evaluation-harness/lm_eval/models/__pycache__/neuralmagic.cpython-310.pyc b/lm-evaluation-harness/lm_eval/models/__pycache__/neuralmagic.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bcddb65015e52d086df10d71328680d2bec7f1cd Binary files /dev/null and b/lm-evaluation-harness/lm_eval/models/__pycache__/neuralmagic.cpython-310.pyc differ diff --git a/lm-evaluation-harness/lm_eval/models/__pycache__/neuralmagic.cpython-311.pyc b/lm-evaluation-harness/lm_eval/models/__pycache__/neuralmagic.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..68ba6bb825bd9f4a5502122af8a61ecc4e380b0e Binary files /dev/null and b/lm-evaluation-harness/lm_eval/models/__pycache__/neuralmagic.cpython-311.pyc differ diff --git a/lm-evaluation-harness/lm_eval/models/__pycache__/neuron_optimum.cpython-311.pyc b/lm-evaluation-harness/lm_eval/models/__pycache__/neuron_optimum.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..653bfb7853b7efb368ef58f420c500bb7cc79ec4 Binary files /dev/null and b/lm-evaluation-harness/lm_eval/models/__pycache__/neuron_optimum.cpython-311.pyc differ diff --git a/lm-evaluation-harness/lm_eval/models/__pycache__/optimum_ipex.cpython-310.pyc b/lm-evaluation-harness/lm_eval/models/__pycache__/optimum_ipex.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..214bbe8c7588ca08824fbb4476a41979eda6f3b4 Binary files /dev/null and b/lm-evaluation-harness/lm_eval/models/__pycache__/optimum_ipex.cpython-310.pyc differ diff --git a/lm-evaluation-harness/lm_eval/models/__pycache__/optimum_ipex.cpython-311.pyc b/lm-evaluation-harness/lm_eval/models/__pycache__/optimum_ipex.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..55ed2e3585c061218797c439644dd14e0a31c647 Binary files /dev/null and b/lm-evaluation-harness/lm_eval/models/__pycache__/optimum_ipex.cpython-311.pyc differ diff --git a/lm-evaluation-harness/lm_eval/models/__pycache__/sglang_causallms.cpython-310.pyc b/lm-evaluation-harness/lm_eval/models/__pycache__/sglang_causallms.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dd921cbf29d48af5b08d121ab33a7b4a749f2552 Binary files /dev/null and b/lm-evaluation-harness/lm_eval/models/__pycache__/sglang_causallms.cpython-310.pyc differ diff --git a/lm-evaluation-harness/lm_eval/models/__pycache__/sglang_generate_API.cpython-310.pyc b/lm-evaluation-harness/lm_eval/models/__pycache__/sglang_generate_API.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8c5d0cb82339405092003b564ba769c51a2eb420 Binary files /dev/null and b/lm-evaluation-harness/lm_eval/models/__pycache__/sglang_generate_API.cpython-310.pyc differ diff --git a/lm-evaluation-harness/lm_eval/models/__pycache__/sglang_generate_API.cpython-311.pyc b/lm-evaluation-harness/lm_eval/models/__pycache__/sglang_generate_API.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..42b272876083f3e20042a8453161005343152ec5 Binary files /dev/null and b/lm-evaluation-harness/lm_eval/models/__pycache__/sglang_generate_API.cpython-311.pyc differ diff --git a/lm-evaluation-harness/lm_eval/models/__pycache__/textsynth.cpython-310.pyc b/lm-evaluation-harness/lm_eval/models/__pycache__/textsynth.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5e9849cd95be6ef08a77b45bd8f5337beb789261 Binary files /dev/null and b/lm-evaluation-harness/lm_eval/models/__pycache__/textsynth.cpython-310.pyc differ diff --git a/lm-evaluation-harness/lm_eval/models/__pycache__/textsynth.cpython-311.pyc b/lm-evaluation-harness/lm_eval/models/__pycache__/textsynth.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0254e1b4f5f9166f48c0290fb393ae874f0184aa Binary files /dev/null and b/lm-evaluation-harness/lm_eval/models/__pycache__/textsynth.cpython-311.pyc differ diff --git a/lm-evaluation-harness/lm_eval/models/__pycache__/utils.cpython-310.pyc b/lm-evaluation-harness/lm_eval/models/__pycache__/utils.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e98584182fb33693d174dc4f39a4ecca14cb53b7 Binary files /dev/null and b/lm-evaluation-harness/lm_eval/models/__pycache__/utils.cpython-310.pyc differ diff --git a/lm-evaluation-harness/lm_eval/models/__pycache__/utils.cpython-311.pyc b/lm-evaluation-harness/lm_eval/models/__pycache__/utils.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c3d25a1381a3130fd95b3bb11e409a9e2781f28c Binary files /dev/null and b/lm-evaluation-harness/lm_eval/models/__pycache__/utils.cpython-311.pyc differ diff --git a/lm-evaluation-harness/lm_eval/models/__pycache__/vllm_causallms.cpython-311.pyc b/lm-evaluation-harness/lm_eval/models/__pycache__/vllm_causallms.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e2284911e64fe610e8cfd7fa96f7ad493896c8fb Binary files /dev/null and b/lm-evaluation-harness/lm_eval/models/__pycache__/vllm_causallms.cpython-311.pyc differ diff --git a/lm-evaluation-harness/lm_eval/models/__pycache__/vllm_vlms.cpython-310.pyc b/lm-evaluation-harness/lm_eval/models/__pycache__/vllm_vlms.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4710df95359416799f22cfd92feeca22613ce19d Binary files /dev/null and b/lm-evaluation-harness/lm_eval/models/__pycache__/vllm_vlms.cpython-310.pyc differ diff --git a/lm-evaluation-harness/lm_eval/models/__pycache__/vllm_vlms.cpython-311.pyc b/lm-evaluation-harness/lm_eval/models/__pycache__/vllm_vlms.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1ca1aa37839b1272c0ae22b6b3f948ad1b5712de Binary files /dev/null and b/lm-evaluation-harness/lm_eval/models/__pycache__/vllm_vlms.cpython-311.pyc differ diff --git a/lm-evaluation-harness/lm_eval/tasks/aclue/aclue_ancient_literature.yaml b/lm-evaluation-harness/lm_eval/tasks/aclue/aclue_ancient_literature.yaml new file mode 100644 index 0000000000000000000000000000000000000000..641befa3aa1920d8dca1c7007a4fe8cd24ab8e77 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/aclue/aclue_ancient_literature.yaml @@ -0,0 +1,4 @@ +"dataset_name": "ancient_literature" +"description": "以下是关于古代文学知识的单项选择题，请直接给出正确答案的选项。\n\n" +"include": "_default_template_yaml" +"task": "aclue_ancient_literature" diff --git a/lm-evaluation-harness/lm_eval/tasks/aclue/aclue_ancient_medical.yaml b/lm-evaluation-harness/lm_eval/tasks/aclue/aclue_ancient_medical.yaml new file mode 100644 index 0000000000000000000000000000000000000000..bceaa702c53a1526fc84cf8f5141570352581a44 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/aclue/aclue_ancient_medical.yaml @@ -0,0 +1,4 @@ +"dataset_name": "ancient_medical" +"description": "以下是关于医古文的单项选择题，请直接给出正确答案的选项。\n\n" +"include": "_default_template_yaml" +"task": "aclue_ancient_medical" diff --git a/lm-evaluation-harness/lm_eval/tasks/aclue/aclue_ancient_phonetics.yaml b/lm-evaluation-harness/lm_eval/tasks/aclue/aclue_ancient_phonetics.yaml new file mode 100644 index 0000000000000000000000000000000000000000..2fe908e531a07466a66f58f2f5009d5111d5a02d --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/aclue/aclue_ancient_phonetics.yaml @@ -0,0 +1,4 @@ +"dataset_name": "ancient_phonetics" +"description": "以下是关于古音学的单项选择题，请直接给出正确答案的选项。\n\n" +"include": "_default_template_yaml" +"task": "aclue_ancient_phonetics" diff --git a/lm-evaluation-harness/lm_eval/tasks/aclue/aclue_couplet_prediction.yaml b/lm-evaluation-harness/lm_eval/tasks/aclue/aclue_couplet_prediction.yaml new file mode 100644 index 0000000000000000000000000000000000000000..63124eed8eb2c2987e7145ee4633e010407641be --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/aclue/aclue_couplet_prediction.yaml @@ -0,0 +1,4 @@ +"dataset_name": "couplet_prediction" +"description": "以下是关于对联的单项选择题，请直接给出正确答案的选项。\n\n" +"include": "_default_template_yaml" +"task": "aclue_couplet_prediction" diff --git a/lm-evaluation-harness/lm_eval/tasks/aclue/aclue_poetry_quality_assessment.yaml b/lm-evaluation-harness/lm_eval/tasks/aclue/aclue_poetry_quality_assessment.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a7a7bee2c4ca59e0dc7b2f3fdc08371a9a585d42 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/aclue/aclue_poetry_quality_assessment.yaml @@ -0,0 +1,4 @@ +"dataset_name": "poetry_quality_assessment" +"description": "以下是关于古诗词质量评估的单项选择题，请直接给出正确答案的选项。\n\n" +"include": "_default_template_yaml" +"task": "aclue_poetry_quality_assessment" diff --git a/lm-evaluation-harness/lm_eval/tasks/aclue/aclue_poetry_sentiment_analysis.yaml b/lm-evaluation-harness/lm_eval/tasks/aclue/aclue_poetry_sentiment_analysis.yaml new file mode 100644 index 0000000000000000000000000000000000000000..6e1367f8043d7e1e9ebcd01dfbaacfbdeb0f9fec --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/aclue/aclue_poetry_sentiment_analysis.yaml @@ -0,0 +1,4 @@ +"dataset_name": "poetry_sentiment_analysis" +"description": "以下是关于诗词情感分类的单项选择题，请直接给出正确答案的选项。\n\n" +"include": "_default_template_yaml" +"task": "aclue_poetry_sentiment_analysis" diff --git a/lm-evaluation-harness/lm_eval/tasks/acpbench/boolq_cot_2shot/act_reach.yaml b/lm-evaluation-harness/lm_eval/tasks/acpbench/boolq_cot_2shot/act_reach.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8c0303ec1e2c6955d2c01c0be7c96e2474686a38 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/acpbench/boolq_cot_2shot/act_reach.yaml @@ -0,0 +1,12 @@ +task: acp_areach_bool +dataset_name: acp_areach_bool +include: _boolq_cot_2shot_yaml +fewshot_config: + sampler: first_n + samples: + - context: 'This is a ferry domain, where the task is to transport cars from their start to their goal locations, using a ferry. Each location is accessible by ferry from each other location. The cars can be debarked or boarded, and the ferry can carry only one car at a time. There are 2 locations and 5 cars, numbered consecutively. Currently, the ferry is at l1, with the car c4 on board. The cars are at locations as follows: c0 and c3 are at l1; c1 and c2 are at l0.' + question: 'Is it possible to transition to a state where the action "travel by sea from location l0 to location l1" can be applied?' + answer: "Let's think step by step. Step 1: Verify if there is a sequence of actions which transforms the current state into a state where the precondition of the action \"travel by sea from location l0 to location l1\" hold. Step 2: The following sequence of actions would transition to such a state: sail from location l1 to location l0, unload the car c4 from the ferry to location l0, board car c1 at location l0. **Final Answer**: Yes." + - context: 'There are several cities, each containing several locations, some of which are airports. There are also trucks, which can drive within a single city, and airplanes, which can fly between airports. The goal is to get some packages from various locations to various new locations. There are 2 trucks and 1 airplane, as well as 4 packages. There are 6 locations across 2 cities. The locations are in cities as follows: l0-0, l0-1, and l0-2 are in c0; l1-1, l1-2, and l1-0 are in c1. Currently, a0 is at l1-0, t1 is at l1-1, t0 is at l0-0, p2 and p1 are in t1, p0 and p3 are in a0.' + question: 'Is it possible to transition to a state where the action "offload the object p0 from the truck p0 at location p1" can be applied?' + answer: "Let's think step by step. Step 1: Verify if there is a sequence of actions which transforms the current state into a state where the precondition of the action \"offload the object p0 from the truck p0 at location p1\" hold. Step 2: Action preconditions are \"p0 is in p0 and p0 is at p1\". Step 3: These facts are not reachable together, as they include mutually exclusive facts \"p0 is in p0 and p0 is at p1\". **Final Answer**: No." diff --git a/lm-evaluation-harness/lm_eval/tasks/acpbench/boolq_cot_2shot/reach.yaml b/lm-evaluation-harness/lm_eval/tasks/acpbench/boolq_cot_2shot/reach.yaml new file mode 100644 index 0000000000000000000000000000000000000000..3a1f06f677c2e84d2231633b9dafc9a7a49358d2 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/acpbench/boolq_cot_2shot/reach.yaml @@ -0,0 +1,13 @@ +task: acp_reach_bool +dataset_name: acp_reach_bool +include: _boolq_cot_2shot_yaml +fewshot_config: + sampler: first_n + samples: + - context: 'This is a ferry domain, where the task is to transport cars from their start to their goal locations, using a ferry. Each location is accessible by ferry from each other location. The cars can be debarked or boarded, and the ferry can carry only one car at a time. There are 2 locations and 2 cars, numbered consecutively. Currently, the ferry is at l1 location and it is empty. The cars are at locations as follows: c0 is at l1; c1 is at l0. is at l1 location and it is empty. The cars are at locations as follows: c0 is at l1; c1 is at l0.' + question: 'Is it possible to transition to a state where the following holds: The ferry is empty and The ferry is at c1 location.' + answer: "Let's think step by step. Step 1: Verify if the following fact(s) hold in current state or if there is a sequence of actions which transforms the current state into a state where they hold: The ferry is empty and The ferry is at c1 location. Step 2: These facts do not hold in the current state. Step 3: The fact There are no cars on the ferry and The ferry is at c1 location is not reachable even by a simple iterative procedure that accumulates all facts made true by applicable actions. **Final Answer**: No." + - context: 'There are several cities, each containing several locations, some of which are airports. There are also trucks, which can drive within a single city, and airplanes, which can fly between airports. The goal is to get some packages from various locations to various new locations. There are 5 trucks and 1 airplane, as well as 4 packages. There are 15 locations across 5 cities. The locations are in cities as follows: l4-2, l4-1, and l4-0 are in c4; l3-2, l3-0, and l3-1 are in c3; l0-2, l0-0, and l0-1 are in c0; l1-0, l1-1, and l1-2 are in c1; l2-1, l2-0, and l2-2 are in c2. Currently, p3 and t1 are at l1-1, t3 is at l3-0, p0 is at l1-0, t4 is at l4-2, p1 and t2 are at l2-0, a0 is at l4-0, t0 is at l0-2, p2 is in a0.' + question: 'Is it possible to transition to a state where the following holds: l2-2 is in l3-2?' + answer: "Let's think step by step. Step 1: Verify if the following fact(s) hold in current state or if there + is a sequence of actions which transforms the current state into a state where they hold: l2-2 is in l3-2. Step 2: These facts do not hold in the current state. Step 3: The fact l2-2 is in l3-2 is not reachable even by a simple iterative procedure that accumulates all facts made true by applicable actions. **Final Answer**: No." diff --git a/lm-evaluation-harness/lm_eval/tasks/acpbench/gen_2shot/next_act.yaml b/lm-evaluation-harness/lm_eval/tasks/acpbench/gen_2shot/next_act.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a264d6449b6630ff008bc5de94e47692dc63a2c8 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/acpbench/gen_2shot/next_act.yaml @@ -0,0 +1,19 @@ +task: acp_nexta_gen +dataset_name: acp_nexta_gen +include: _gen_yaml_2shot +fewshot_config: + sampler: first_n + samples: + - context: "A robot is in a grid and can only move to places that are connected to its current position. The grid size is 5x5, and the locations are of the form fi-jf (e.g., f3-2f or f0-1f). The grid cells are connected to their neighbors (e.g., f1-2f is connected to the four neighbors f0-2f, f2-2f, f1-1f, and f1-3f). Some positions on the grid are locked and can be opened with a key of a matching shape. The robot has an arm that can pick up a key when the key is in same location as the robot and the arm is empty. There are 2 keys in 1 different shapes: Key key0-1 is of shape shape0, Key key0-0 is of shape shape0. Currently, the robot is at position f4-0f and its arm is empty. All the positions are open except the following: f4-2f has shape0 shaped lock. Key key0-0 is at position f3-0f. Key key0-1 is at position f1-3f. The goal is to reach a state where the following facts hold: Key key0-0 is at f2-0f location and Key key0-1 is at f1-3f location. The available actions are: (unlock ?curpos ?lockpos ?key ?shape) - unlock place ?lockpos with key ?key of shape ?shape from current position place ?curpos, (move ?curpos ?nextpos) - travel from the current position ?curpos to the next position ?nextpos, (pickup ?curpos ?key) - pick up key ?key from place ?curpos, (pickup-and-loose ?curpos ?newkey ?oldkey) - pick up the key ?newkey at the current position place ?curpos and loose the key ?oldkey being held, and (putdown ?curpos ?key) - put down the key ?key at the current position ?curpos." + question: "What is the next action that takes us towards the goal?" + answer: "(move f4-0f f3-0f)" + - context: "There are several cities, each containing several locations, some of which are airports. There are also trucks, which can drive within a single city, and airplanes, which can fly between airports. The goal is to get some packages from various locations to various new locations. There are 2 trucks and 1 airplane, as well as 4 packages. There are 4 locations across 2 cities. The locations are in cities as follows: l0-1 and l0-0 are in c0; l1-1 and l1-0 are in c1. Currently, t0 is at l0-1, a0 is at l0-0, t1 and p1 are at l1-0, p2, p0, and p3 are in t1. The goal is to reach a state where the following facts hold: p3 is at l0-1, p2 is at l1-0, p1 is at l1-0, and p0 is at l0-0. The available actions are: (load-truck ?obj ?truck ?loc) - load object ?obj into truck ?truck at location ?loc, (load-airplane ?obj ?airplane ?loc) - load the object ?obj from location ?loc onto the airplane ?airplane, (unload-truck ?obj ?truck ?loc) - unload the object ?obj from the truck ?truck at location ?loc, (unload-airplane ?obj ?airplane ?loc) - unload object ?obj from airplane ?airplane at location ?loc, (drive-truck ?truck ?loc-from ?loc-to ?city) - drive the truck ?truck in city ?city from location ?loc-from to location ?loc-to, and (fly-airplane ?airplane ?loc-from ?loc-to) - operate the airplane ?airplane from airport ?loc-from to airport ?loc-to." + question: "What is the next action that takes us towards the goal?" + answer: "(drive-truck t0 l0-1 l0-0 c0)" +doc_to_text: "**Question**: {{context}} {{question}} Each action starts with an opening parenthesis and ends with closing parenthesis. Provide only the action. **Final Answer**:" +filter_list: + - name: "acp_grammar_parse" + filter: + - function: "ACP_grammar_filter" + grammar_task: "action_name" + - function: "take_first" diff --git a/lm-evaluation-harness/lm_eval/tasks/acpbench/gen_2shot/prog.yaml b/lm-evaluation-harness/lm_eval/tasks/acpbench/gen_2shot/prog.yaml new file mode 100644 index 0000000000000000000000000000000000000000..6267f29acdd0c9767aa11277dadfa35d763bfeb1 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/acpbench/gen_2shot/prog.yaml @@ -0,0 +1,20 @@ +task: acp_prog_gen +dataset_name: acp_prog_gen +include: _gen_yaml_2shot +fewshot_config: + sampler: first_n + samples: + - context: "A robot is in a grid and can only move to places that are connected to its current position. \nThe grid size is 5x5, and the locations are of the form fi-jf (e.g., f3-2f or f0-1f). The grid cells are connected to their neighbors (e.g., f1-2f is connected to the four neighbors f0-2f, f2-2f, f1-1f, and f1-3f). Some positions on the grid are locked and can be opened with a key of a matching shape. The robot has an arm that can pick up a key when the key is in same location as the robot and the arm is empty. \nThere are 2 keys in 0 different shapes: Key key0-0 is of shape shape0, Key key0-1 is of shape shape0. \nCurrently, the robot is at position f0-1f and its arm is empty. All the positions are open except the following: f4-2f has shape0 shaped lock. Key key0-1 is at position f1-3f. Key key0-0 is at position f0-1f. The available propositions are: (at ?r ?x) - Key ?r is at ?x location, (at-robot ?x) - Robot is at ?x location, (locked ?x) - Location ?x is locked, (holding ?k) - Robot is holding ?k, (open ?x) - Location ?x is open, and (arm-empty) - Robot's arm is empty." + question: "Break down the outcomes of performing the action \"retrieve the key key0-0 from its current position f0-1f\" into two lists, positive effects and negative effects. Positive effects are the propositions that are false in the current state but will become true after performing the action. Negative effects are the propositions that are true in the current state and will become false after performing the action." + answer: "[(holding key0-0)] [(arm-empty), (at key0-0 f0-1f)]" + - context: "There are several cities, each containing several locations, some of which are airports. There are also trucks, which can drive within a single city, and airplanes, which can fly between airports. The goal is to get some packages from various locations to various new locations. There are 2 trucks and 1 airplane, as well as 4 packages. There are 4 locations across 2 cities. The locations are in cities as follows: l1-1 and l1-0 are in c1; l0-1 and l0-0 are in c0. Currently, p2, t1, p1, p3, a0, and p0 are at l1-0, t0 is at l0-1. The available propositions are: (at ?obj ?loc) - ?obj is at ?loc and (in ?obj1 ?obj2) - ?obj1 is in ?obj2." + question: "Break down the outcomes of performing the action \"load object p3 into truck t1 at location l1-0\" into two lists, positive effects and negative effects. Positive effects are the propositions that are false in the current state but will become true after performing the action. Negative effects are the propositions that are true in the current state and will become false after performing the action." + answer: "[(in p3 t1)] [(at p3 l1-0)]" +doc_to_text: "**Question**: {{context}} {{question}} Provide only the two lists with the ground propositions. **Final Answer**:" +filter_list: + - name: "acp_grammar_parse" + filter: + - function: "ACP_grammar_filter" + grammar_task: "progression_list" + clean: "pos_neg" + - function: "take_first" diff --git a/lm-evaluation-harness/lm_eval/tasks/acpbench/gen_2shot/val.yaml b/lm-evaluation-harness/lm_eval/tasks/acpbench/gen_2shot/val.yaml new file mode 100644 index 0000000000000000000000000000000000000000..5dc02acf6a8c4948305e8f827c99064ba6b440a5 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/acpbench/gen_2shot/val.yaml @@ -0,0 +1,19 @@ +task: acp_val_gen +dataset_name: acp_val_gen +include: _gen_yaml_2shot +fewshot_config: + sampler: first_n + samples: + - context: "A robot is in a grid and can only move to places that are connected to its current position. The grid size is 5x5, and the locations are of the form fi-jf (e.g., f3-2f or f0-1f). The grid cells are connected to their neighbors (e.g., f1-2f is connected to the four neighbors f0-2f, f2-2f, f1-1f, and f1-3f). Some positions on the grid are locked and can be opened with a key of a matching shape. The robot has an arm that can pick up a key when the key is in same location as the robot and the arm is empty. There are 2 keys in 1 different shapes: Key key0-0 is of shape shape0, Key key0-1 is of shape shape0. Currently, the robot is at position f3-3f and its arm is empty. All the positions are open except the following: f2-0f has shape0 shaped lock, f4-2f has shape0 shaped lock. Key key0-1 is at position f1-3f. Key key0-0 is at position f2-2f. The goal is to reach a state where the following facts hold: Key key0-0 is at f2-0f location and Key key0-1 is at f1-3f location. The available actions are: (unlock ?curpos ?lockpos ?key ?shape) - unlock the place ?lockpos with the key ?key of the shape ?shape from the current position place ?curpos, (move ?curpos ?nextpos) - travel from the current position ?curpos to the next position ?nextpos, (pickup ?curpos ?key) - pick up key ?key from place ?curpos, (pickup-and-loose ?curpos ?newkey ?oldkey) - pick up the key ?newkey from the current position ?curpos and loose the key ?oldkey which is being held, and (putdown ?curpos ?key) - put down key ?key at current position place ?curpos." + question: "What is the first inapplicable action in the next sequence of actions: [(move f3-3f f3-2f), (move f3-2f f2-2f), (pickup f2-2f key0-0), (pickup-and-loose f4-0f key0-0 key0-1), (unlock f2-1f f2-0f key0-0 shape0), (move f2-1f f2-0f), (putdown f2-0f key0-0), (move f2-0f f2-1f)]?" + answer: "3" + - context: "There are several cities, each containing several locations, some of which are airports. There are also trucks, which can drive within a single city, and airplanes, which can fly between airports. The goal is to get some packages from various locations to various new locations. There are 2 trucks and 1 airplane, as well as 4 packages. There are 4 locations across 2 cities. The locations are in cities as follows: l0-1 and l0-0 are in c0; l1-1 and l1-0 are in c1. Currently, t1 and p0 are at l1-1, t0 is at l0-1, p3, p2, and p1 are at l1-0, a0 is at l0-0. The goal is to reach a state where the following facts hold: p2 is at l1-0, p3 is at l0-1, p0 is at l0-0, and p1 is at l1-0. The available actions are: (load-truck ?obj ?truck ?loc) - load object ?obj into truck ?truck at location ?loc, (load-airplane ?obj ?airplane ?loc) - load the object ?obj from location ?loc onto the airplane ?airplane, (unload-truck ?obj ?truck ?loc) - unload the object ?obj from the truck ?truck at location ?loc, (unload-airplane ?obj ?airplane ?loc) - unload object ?obj from airplane ?airplane at location ?loc, (drive-truck ?truck ?loc-from ?loc-to ?city) - navigate the truck ?truck from its current location ?loc-from in city ?city to the new location ?loc-to within the same city, and (fly-airplane ?airplane ?loc-from ?loc-to) - fly the airplane ?airplane from location ?loc-from to location ?loc-to." + question: "What is the first inapplicable action in the next sequence of actions: [(load-truck p0 t1 l1-1), (drive-truck t1 l1-1 l1-0 c1), (unload-truck p0 t1 l1-0), (fly-airplane a0 l0-0 l1-0), (unload-truck p3 t0 l0-1), (load-airplane p3 a0 l1-0), (fly-airplane a0 l1-0 l0-0), (unload-airplane p0 a0 l0-0), (unload-airplane p3 a0 l0-0), (drive-truck t0 l0-1 l0-0 c0), (load-truck p3 t0 l0-0), (drive-truck t0 l0-0 l0-1 c0), (unload-truck p3 t0 l0-1)]?" + answer: "4" +doc_to_text: "**Question**: {{context}} {{question}} Provide only the index of the action. **Final Answer**:" +filter_list: + - name: "acp_grammar_parse" + filter: + - function: "ACP_grammar_filter" + grammar_task: "index" + - function: "take_first" diff --git a/lm-evaluation-harness/lm_eval/tasks/acpbench/gen_2shot_with_pddl/acp_grammar.lark b/lm-evaluation-harness/lm_eval/tasks/acpbench/gen_2shot_with_pddl/acp_grammar.lark new file mode 100644 index 0000000000000000000000000000000000000000..036bd675faacb044ca5bb2ef66b3dfac47943815 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/acpbench/gen_2shot_with_pddl/acp_grammar.lark @@ -0,0 +1,23 @@ +NAME: /[a-zA-Z][a-zA-Z0-9-_]*/ +LPAR : "(" +RPAR : ")" +LSPAR: "[" +RSPAR: "]" +COMMA: "," +WS: /[ \n]/ + +action_none : "None" + +action_name : LPAR NAME (WS NAME)* RPAR + +action_list : (action_name WS?)* + +prog_list : action_name* (COMMA action_name)* + +progression_list : LSPAR prog_list RSPAR LSPAR prog_list RSPAR + +act : action_name | action_none + +index: /[0-9]+[0-9]*/ + +start: action_list diff --git a/lm-evaluation-harness/lm_eval/tasks/acpbench/gen_2shot_with_pddl/acp_utils.py b/lm-evaluation-harness/lm_eval/tasks/acpbench/gen_2shot_with_pddl/acp_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..5051b68cbf7b5ef384f2ec498f2759409383c7b7 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/acpbench/gen_2shot_with_pddl/acp_utils.py @@ -0,0 +1,1128 @@ +import json +import os +from abc import ABC, abstractmethod +from collections import defaultdict +from pathlib import Path + +from lm_eval.api.registry import register_filter +from lm_eval.filters.extraction import RegexFilter + + +try: + import tempfile + + import tarski + from kstar_planner import planners as kp + from lark import Lark + from lark.lexer import Token + from lark.visitors import Visitor + from pddl.core import Problem + from pddl.parser.domain import DomainParser + from pddl.parser.problem import ProblemParser + from tarski.grounding.common import StateVariableLite + from tarski.grounding.lp_grounding import LPGroundingStrategy + from tarski.io import PDDLReader + from tarski.io import fstrips as iofs + from tarski.syntax.formulas import is_atom + from tarski.syntax.transform.action_grounding import ( + ground_schema_into_plain_operator_from_grounding, + ) + from tarski.util import SymbolIndex +except ModuleNotFoundError: + raise ModuleNotFoundError( + "`lark>=1.1.9`, `tarski[clingo]==0.8.2`, `pddl==0.4.2` and `kstar-planner==1.4.2` are required for evaluating the generative tasks. \ +Please install via pip install lm-eval[acpbench] or pip install -e .[acpbench]", + ) + + +######################################################################### +# Grammar + + +GRAMMAR_FILE = os.path.join( + os.path.dirname(os.path.abspath(__file__)), "acp_grammar.lark" +) + + +class ACPBench_Visitor(Visitor): + def __init__(self) -> None: + super().__init__() + self.action_lists = None + self.action_names = None + self.progression_lists = None + self.prog_lists = None + self.indexes = None + + def action_list(self, tree): + self.action_lists = [] + + def prog_list(self, tree): + if self.prog_lists is not None: + self.progression_lists.append(self.prog_lists) + self.prog_lists = [] + + def progression_list(self, tree): + self.progression_lists = [] + + def action_none(self, tree): + self.action_names = "None" + + def action_name(self, tree): + act_name = "(" + "".join(tree.children[1:-1]) + ")" + self.action_names = act_name + if self.action_lists is not None: + self.action_lists.append(act_name) + if self.prog_lists is not None: + self.prog_lists.append(act_name) + + def index(self, tree): + self.indexes = "".join(tree.children) + if not self.indexes.isnumeric(): + self.indexes = None + + +class ACPGrammarParser(object): + def __init__(self, task) -> None: + self.task = task + with open(GRAMMAR_FILE) as f: + grammar = f.read() + self.acp_parser = Lark(grammar, start=task, parser="lalr") + + def parse(self, input, debug=False): + def ignore_errors(e): + if hasattr(e, "token") and e.token.type == "$END": + for x in e.expected: + if x != "WS": + e.interactive_parser.feed_token( + Token(x, self.acp_parser.get_terminal(x).pattern.value) + ) + + return True + + input = input.replace("\n", "") + input = input.strip() + try: + tree = self.acp_parser.parse(input, on_error=ignore_errors) + + if debug: + print(tree) + visitor = ACPBench_Visitor() + visitor.visit_topdown(tree) + if self.task == "action_list": + return visitor.action_lists + elif self.task == "act": + return visitor.action_names + elif self.task == "action_name": + return visitor.action_names + elif self.task == "index": + return visitor.indexes + elif self.task == "progression_list": + if visitor.prog_lists not in visitor.progression_lists: + visitor.progression_lists.append(visitor.prog_lists) + return visitor.progression_lists + except Exception as e: + if debug: + print("exception") + print(e) + return None + + +############################################################################## +# Utils + + +# Used in next action +def is_on_optimal_plan(domain, problem, action, opt): + with ( + tempfile.NamedTemporaryFile() as domain_temp, + tempfile.NamedTemporaryFile() as problem_temp, + ): + with open(str(domain_temp.name), "w", encoding="utf8") as file: + file.write(domain.lower()) + with open(str(problem_temp.name), "w", encoding="utf8") as file: + file.write(problem.lower()) + + # Here, we need to keep the temp files live until the end of the function + try: + P = STRIPS(str(domain_temp.name), str(problem_temp.name)) + except Exception: + # Unsolvable + return False + + a = P.get_action_or_none(action[1:-1]) + if a is None: + return False + state = P.init + next_state = progress(state, a) + if opt is None: + # Get an optimal plan cost + plans = generate_optimal_plans_for_problem_state( + P, state, num_plans=1, timeout=5 + ) + opt = len(plans[0]["actions"]) + else: + opt = int(opt) + + # Getting an optimal plan for the next state + next_plans = generate_optimal_plans_for_problem_state( + P, next_state, num_plans=1, timeout=5 + ) + if next_plans is None: + return False + next_opt = len(next_plans[0]["actions"]) + return next_opt + 1 == opt + + +# Used in justification +def is_plan(domain, problem, new_plan): + P = get_STRIPS(domain, problem) + if P is None: + # Unsolvable + return False + + # Check if new_plan is a plan + current_state = P.init + for action in new_plan: + applicable_actions = P.get_applicable_actions(current_state) + app_actions_list = [f"({a.name.lower()})" for a in applicable_actions] + if action.lower() not in app_actions_list: + return False + a = applicable_actions[app_actions_list.index(action.lower())] + current_state = progress(current_state, a) + return entails(current_state, P.goal) + + +# Used in action reachability +def get_action_preconditions(domain, problem, action): + P = get_STRIPS(domain, problem) + + assert P is not None, f"Domain\n{domain}\nProblem\n{problem}\nAction: {action}" + a = P.get_action_or_none(action[1:-1]) + if a is None: + return a + + return [f"({f})" for f in a.pres] + + +def generate_optimal_plans_for_problem_state(P, state, num_plans, timeout): + import tempfile + + with ( + tempfile.NamedTemporaryFile() as domain_temp, + tempfile.NamedTemporaryFile() as problem_temp, + ): + create_tmp_dom_prob_replace_init(P, state, domain_temp, problem_temp) + plans = generate_top_q_plans( + domain=str(domain_temp.name), + problem=str(problem_temp.name), + num_plans=num_plans, + quality_bound=1.0, + timeout=timeout, + ) + # print(plans) + if plans is None or len(plans["plans"]) == 0: + return None + return plans["plans"] + + +def generate_top_q_plans(domain, problem, num_plans=10, quality_bound=1.0, timeout=30): + # print("Running K* planner") + plans = kp.plan_unordered_topq( + domain_file=Path(domain), + problem_file=Path(problem), + number_of_plans_bound=num_plans, + quality_bound=quality_bound, + timeout=timeout, + ) + return plans + + +# Used in (action) reachability +def is_unsolvable_new_goal(domain, problem, new_goal): + goal = extract_goal(problem) + new_problem = problem.replace(goal, f"(:goal {new_goal} )") + return is_unsolvable(domain, new_problem) + + +def is_unsolvable(domain, problem): + with ( + tempfile.NamedTemporaryFile() as domain_temp, + tempfile.NamedTemporaryFile() as problem_temp, + ): + with open(str(domain_temp.name), "w", encoding="utf8") as file: + file.write(str(domain)) + with open(str(problem_temp.name), "w", encoding="utf8") as file: + file.write(str(problem)) + + plans = kp.plan_unordered_topq( + domain_file=Path(str(domain_temp.name)), + problem_file=Path(str(problem_temp.name)), + quality_bound=1.0, + number_of_plans_bound=1, + timeout=3, + ) + + if len(plans["planner_error"]) > 0: + fl = plans["planner_error"].split("\n")[0] + print(f"Planner error: {fl}") + return False + if plans is None or len(plans["plans"]) == 0: + return plans["unsolvable"] + return False + + +def extract_goal(prob): + a = prob.split("(:goal")[1] + cp = 1 + for i, c in enumerate(a): + if c == ")": + cp -= 1 + if c == "(": + cp += 1 + if cp == 0: + return "(:goal" + a[: i + 1] + + assert False + + +def entails(state, partialstate): + return partialstate <= state + + +def progress(state, act): + assert entails(state, act.pres), ( + "Cannot progress with inconsistent state / action precondition:\n\t Action: " + + act.name + + "\n\t State: \n\t\t" + + "\n\t\t".join(state) + ) + return (state - act.dels) | act.adds + + +def regress(state, act): + assert len(state & act.dels) == 0, ( + "Cannot regress with inconsistent state / action delete effect:\n\t Action: " + + act.name + + "\n\t State: \n\t\t" + + "\n\t\t".join(state) + ) + return (state - act.adds) | act.pres + + +def get_STRIPS(domain, problem): + with ( + tempfile.NamedTemporaryFile() as domain_temp, + tempfile.NamedTemporaryFile() as problem_temp, + ): + with open(str(domain_temp.name), "w", encoding="utf8") as file: + file.write(domain.lower()) + with open(str(problem_temp.name), "w", encoding="utf8") as file: + file.write(problem.lower()) + + try: + P = STRIPS(str(domain_temp.name), str(problem_temp.name)) + return P + except Exception as e: + print(f"||{e}||") + return None + + +def create_tmp_dom_prob_replace_init(P, state, result_domain_file, result_problem_file): + d, p = P.PDDL_replace_init_pddl_parser(state) + with open(str(result_domain_file.name), "w", encoding="utf8") as file: + file.write(str(d)) + with open(str(result_problem_file.name), "w", encoding="utf8") as file: + file.write(str(p)) + + return d, p + + +def fix_name(s): + # (act param) + if "(" == s[0] and ")" == s[-1]: + return s[1:-1] + # make it space separated + s = s.replace(", ", " ").replace(",", " ") + # act(param) + if "(" in s: + assert ")" == s[-1], f"Broken name? {s}" + s = s.replace("(", " ").replace(")", "") + # act param + return s + + +def get_atoms_pddl(d, p, atoms): + objs = set() + preds = defaultdict(list) + for atom in atoms: + a = atom.lower().strip().split(" ") + args = a[1:] + preds[a[0]].append(args) + objs |= set(args) + + constants = [o for o in p.objects | d.constants if o.name.lower() in objs] + constants_dict = {} + for c in constants: + constants_dict[c.name.lower()] = c + assert len(objs) == len(constants), ( + f"Could not identify all objects: {objs - set(constants_dict.keys())} not found, {set(constants_dict.keys()) - objs} should not be there" + ) + + state = [] + covered_preds = set() + for f in d.predicates: + name = f.name.lower() + if name in preds: + covered_preds.add(name) + assert len(preds[name][0]) == f.arity, ( + f"The arity does not match: {preds[name]} vs {f.terms}" + ) + # Going over the lists of objects, adding ground predicate for each + for ob in preds[name]: + c = [constants_dict[o] for o in ob] + state.append(f(*c)) + assert len(covered_preds) == len(preds.keys()), ( + f"Covered predicates: \n{sorted(list(covered_preds))} vs \n{sorted(list(preds.keys()))}" + ) + return set(state) + + +class Action: + def __init__(self, name, pre, add, delete): + self.name = name + self.pres = pre + self.adds = add + self.dels = delete + + def __str__(self): + pres = "{" + ", ".join([f"({a})" for a in self.pres]) + "}" + adds = "{" + ", ".join([f"({a})" for a in self.adds]) + "}" + dels = "{" + ", ".join([f"({a})" for a in self.dels]) + "}" + + return f"< {self.name}, {pres}, {adds}, {dels} >" + + def toJSON(self): + return json.dumps( + { + "name": self.name, + "preconditions": [f"({a})" for a in self.pres], + "add_effects": [f"({a})" for a in self.adds], + "delete_effects": [f"({a})" for a in self.dels], + }, + sort_keys=True, + indent=4, + ) + + def __repr__(self): + return self.name + + def __eq__(self, action): + return self.name == action.name + + def __hash__(self): + return hash(self.name) + + +class STRIPS: + def __init__(self, domain, problem): + self.domain_file = domain + self.problem_file = problem + self.reader = PDDLReader(raise_on_error=True) + self.reader.parse_domain(domain) + self.problem = self.reader.parse_instance(problem) + (self.grounded_fluents, init, goal, self.operators, self.grounder) = ( + self.ground_problem(self.problem) + ) + + self.fluents = set([fix_name(str(f)) for f in self.grounded_fluents]) + self.fluents_map = dict() + for f in self.grounded_fluents: + self.fluents_map[fix_name(str(f))] = f + self.init = set([fix_name(str(f)) for f in init]) + self.goal = set([fix_name(str(f)) for f in goal]) + self.actions = set() + self.action_map = {} + self.init_fluents = [self.fluents_map[f] for f in self.init] + + self.static_predicates = [i.name for i in self.grounder.static_symbols] + for op in self.operators: + act = self.operator_to_action(op) + self.actions.add(act) + self.action_map[act.name.lower()] = act + + def __str__(self): + fluents = "P = {" + ", ".join([f"({a})" for a in self.fluents]) + "}" + init = "I = {" + ", ".join([f"({a})" for a in self.init]) + "}" + goal = "G = {" + ", ".join([f"({a})" for a in self.goal]) + "}" + actions = "A = {" + "\n ".join([a.__str__() for a in self.actions]) + "}" + return fluents + ",\n" + init + "\n" + goal + "\n" + actions + + def toJSON(self): + actions = [a.toJSON() for a in self.actions] + return json.dumps( + { + "fluents": list(self.fluents), + "initial_state": list(self.init), + "goal": list(self.goal), + "actions": actions, + }, + sort_keys=True, + indent=4, + ) + + def operator_to_action(self, op, check_fluents=True, check_static=False): + adds = { + fix_name(str(f.atom)) for f in op.effects if isinstance(f, iofs.AddEffect) + } & self.fluents + dels = { + fix_name(str(f.atom)) for f in op.effects if isinstance(f, iofs.DelEffect) + } & self.fluents + pre = self.fix_pre_name(op.precondition) + if check_fluents: + pre = pre & self.fluents + if check_static: + pre = {p for p in pre if p.split()[0] not in self.static_predicates} + act = Action(fix_name(str(op)), pre, adds, dels) + return act + + def fix_pre_name(self, precondition): + if not is_atom(precondition): + return {fix_name(str(f)) for f in precondition.subformulas} + return {fix_name(str(precondition))} + + def action(self, name): + return self.action_map[fix_name(name).lower()] + + def get_action_or_none(self, name): + if "(" in name and ")" != name[-1]: + return None + return self.action_map.get(fix_name(name).lower(), None) + + def fluent(self, name): + return fix_name(name) + + def static_symbols(self): + return list(self.grounder.static_symbols) + + def fluent_symbols(self): + return list(self.grounder.fluent_symbols) + + def get_grounded_atoms(self, symbol): + variables = SymbolIndex() + lang = symbol.language + key = "atom_" + symbol.name + model = self.grounder._solve_lp() + if ( + key in model + ): # in case there is no reachable ground state variable from that fluent symbol + for binding in model[key]: + binding_with_constants = tuple(lang.get(c) for c in binding) + variables.add(StateVariableLite(symbol, binding_with_constants)) + return variables + + def get_applicable_actions(self, s): + return [a for a in self.actions if entails(s, a.pres)] + + def ground_problem(self, problem): + grounder = LPGroundingStrategy(problem, include_variable_inequalities=True) + action_groundings = grounder.ground_actions() + operators = [] + for action_name, groundings in action_groundings.items(): + action = problem.get_action(action_name) + for grounding in groundings: + operators.append( + ground_schema_into_plain_operator_from_grounding(action, grounding) + ) + + grounded_fluents = set( + [ + grounded_fluent.to_atom() + for grounded_fluent in grounder.ground_state_variables().objects + ] + ) + init = [f for f in problem.init.as_atoms() if f in grounded_fluents] + if isinstance(problem.goal, tarski.syntax.Atom): + goal = [problem.goal] + else: + goal = [f for f in problem.goal.subformulas if f in grounded_fluents] + + return (grounded_fluents, init, goal, operators, grounder) + + def get_static(self): + static_symbols = self.static_symbols() + ret = [] + for symbol in static_symbols: + ret.extend(self.get_grounded_atoms(symbol)) + return set([fix_name(str(x)) for x in ret]) + + def PDDL_replace_init_pddl_parser(self, s): + d = DomainParser()(open(self.domain_file, "r").read().lower()) + p = ProblemParser()(open(self.problem_file, "r").read().lower()) + + new_state = get_atoms_pddl(d, p, s | self.get_static()) + + new_p = Problem( + p.name, domain=d, objects=p.objects, init=new_state, goal=p.goal + ) + + return d, new_p + + +def parse_ans(response: str, parser: ACPGrammarParser, task: str): + return [parser.parse(clean_answer(resp, task)) for resp in response] + + +# def parse_ans(response : str, parser : ACPGrammarParser, task : str): +# ans = [parser.parse(clean_answer(resp, task), debug=True) for resp in response] +# if any(elem is None for elem in ans) or any(elem is None for elem in ans[0]): +# return None +# return ans + + +def remove_garbage(s): + while True: + if s.endswith("."): + s = s[:-1] + elif s.endswith("\n"): + s = s[:-2] + else: + break + return s.rstrip() + + +def compare_str(s1, s2): + return remove_garbage(s1).lower() == remove_garbage(s2).lower() + + +def compare(l1, l2): + if not isinstance(l1, list): + return compare_str(l1, l2) + if not isinstance(l2, list): + return False + for i, v in enumerate(l1): + if not compare(v, l2[i]): + return False + return True + + +def check_prog_response(resp): + if ( + "Positive Effects".lower() in resp.lower() + and "Negative Effects".lower() in resp.lower() + ): + if "[" not in resp: + return True + return False + + +def clean_answer(resp, task): + # Minor cleanup + if "progression_gen" in task: + # Check for Positive Effects and Negative Effects instead of separation + if check_prog_response(resp): + # replace **Positive Effects** with "[" + # replace **Negative Effects** with "] [" + # append "]" to the end + resp2 = resp.lower() + resp2 = resp2.replace("*", "") + resp2 = resp2.replace("positive effects", "[") + resp2 = resp2.replace("negative effects", "] [") + resp2 = resp2 + "]" + return resp2 + if "action_justification_gen" in task: + # Check for "simplified plan:" + if "simplified plan:" in resp.lower(): + resp2 = resp.lower() + resp2 = resp2.replace("*", "") + resp2 = resp2.split("simplified plan:")[1] + return resp2 + return resp + + +def get_grammar_task(task): + # print(task) + if task == "reachable_atom_gen": + return "act" + elif task == "progression_gen": + return "progression_list" + elif task == "validation_gen": + return "index" + elif task == "reachable_action_gen": + return "act" + elif task == "action_justification_gen": + return "action_list" + elif task == "landmarks_gen": + return "act" + elif task == "goal_closer_gen": + return "action_name" + elif task == "applicable_actions_gen": + return "action_list" + + +############################################################################## +# Evaluators + + +def fix_action_name(a): + assert a.startswith("(") and a.endswith(")") + return "(" + " ".join([x.strip() for x in a[1:-1].split(" ") if len(x) > 0]) + ")" + + +def str_remove_before_first_parentheses(s): + if s.startswith("("): + return s + try: + return s[s.index("(") :] + except Exception: + return "" + + +def str_remove_after_last_parentheses(s): + if s.endswith(")"): + return s + + i = s.rfind(")") + + if i == -1: + return "" + return s[: i + 1] + + +def cleanup_answer(ans): + if isinstance(ans, str): + ans = str_remove_before_first_parentheses(ans) + ans = str_remove_after_last_parentheses(ans) + ans = ans.lower() + ans = ( + ans.replace(")\n(", ")######(") + .replace("),(", ")######(") + .replace(") (", ")######(") + .split("######") + ) + return ans + if isinstance(ans, list): + res = [] + for x in ans: + res.extend(cleanup_answer(x)) + return res + + +def set_equal(ans1, ans2): + return set(ans1) == set(ans2) + + +class BaseEvaluator(ABC): + def __init__(self) -> None: + self.scores = [] + + @abstractmethod + def get_score(self, ans, doc): + pass + + def add_scores(self, scores): + self.scores.extend(scores) + + def get_avg_score(self): + avg_score = sum(self.scores) / len(self.scores) + return avg_score + + +def get_evaluator(group): + if group == "applicable_actions_gen": + return ApplicabilityEvaluator() + elif group == "progression_gen": + return ProgressionEvaluator() + elif group == "validation_gen": + return ValidationEvaluator() + elif group == "reachable_atom_gen": + return ReachabilityEvaluator() + elif group == "goal_closer_gen": + return NextActionEvaluator() + elif group == "action_justification_gen": + return JustificationEvaluator() + elif group == "landmarks_gen": + return LandmarksEvaluator() + elif group == "reachable_action_gen": + return ActionReachabilityEvaluator() + assert True, f"Group {group} not found" + + +""" +Action Reachability task: generate a valid action that is not applicable to any reachable state. +answer: A subset of actions that are known to be unreachable (not an exhaustive set). + It is empty only when we *know* that there are no such actions. +""" + + +class ActionReachabilityEvaluator(BaseEvaluator): + def get_score(self, ans, doc): + real_answer = doc["answer"] + if not real_answer or len(real_answer) == 0: + # The correct answer is None + self.add_scores( + ["none" == x.strip().lower() if x is not None else False for x in ans] + ) + else: + for x in ans: + if x is None: + self.scores.append(False) + continue + action = x.strip().lower() + if action in real_answer: + # The answer is in the subset of stored correct answers + self.scores.append(True) + continue + prec = get_action_preconditions( + doc["PDDL_domain"].lower(), doc["PDDL_problem"].lower(), action + ) + if prec is None: + # The answer does not correspond to a valid action + self.scores.append(False) + else: + # Need to run a planner on a task with the answer action preconditions as the new goal + prec = f"(and {' '.join(prec)})" + self.scores.append( + is_unsolvable_new_goal( + doc["PDDL_domain"].lower(), + doc["PDDL_problem"].lower(), + prec, + ) + ) + + return self.get_avg_score() + + +""" +Action Applicability task: generate all actions that are applicable in the current state. +answer: A set of all applicable actions. +""" + + +class ApplicabilityEvaluator(BaseEvaluator): + def get_score(self, ans, doc): + real_answer = doc["answer"] + real_answer = [a.lower() for a in real_answer] + ans = [[fix_action_name(a) for a in x] if x is not None else None for x in ans] + + # Check if the answer is equal (as a set) to the real stored answer + self.add_scores( + [ + set_equal(real_answer, cleanup_answer(x)) if x is not None else False + for x in ans + ] + ) + return self.get_avg_score() + + +def is_subsequence(plan, new_plan): + i = 0 + for a in plan: + if a == new_plan[i]: + i += 1 + if len(new_plan) == i: + # Done + return True + return False + + +def is_subsequence_and_plan(domain, problem, plan, new_plan): + if len(plan) <= len(new_plan): + return False + if not is_subsequence(plan, new_plan): + return False + return is_plan(domain, problem, new_plan) + + +""" +Justification task: generate a proper subsequence of the given plan that is also a plan. +answer: A list of examples of actions that can be removed (ignored in evaluation). +""" + + +class JustificationEvaluator(BaseEvaluator): + def get_score(self, ans, doc): + # Sequence of actions (plan) from the question + if "inputs" in doc: # old field name + seq = doc["inputs"][19:-147] + else: + seq = doc["question"][19:-147] + seq = seq.replace(") (", ")######(").split("######") + for x in ans: + if x is None: + self.scores.append(False) + continue + # An answer plan candidate + x = [fix_action_name(a) for a in x] + if len(x) == 0: + # Wrong answer - never an empty sequence + self.scores.append(0) + continue + # Check if the plan candidate from the answer (a) is a proper subsequence of the plan in the question and (b) is a plan. + self.scores.append( + is_subsequence_and_plan( + doc["PDDL_domain"].lower(), doc["PDDL_problem"].lower(), seq, x + ) + ) + return self.get_avg_score() + + +""" +Landmarks task: generate a fact that is a non-trivial landmark for the current state. +answer: A list of facts that are found to be landmarks and a list of facts that are found to be non-landmarks. + +The questions are generated only for cases where all facts either + (a) hold in the current state, + (b) true in goal, + (c) are found to be landmarks, or + (d) are found to be non-landmarks. +In such cases, the evaluation is simple, it does not require checking whether a fact is a landmark, it was +already done during question generation. +""" + + +class LandmarksEvaluator(BaseEvaluator): + def get_score(self, ans, doc): + # The set of facts that are found to be landmarks + real_answer = doc["answer"] + real_answer_yes = [a.lower() for a in real_answer["yes"]] + + for x in ans: + if x is None: + self.scores.append(False) + continue + if x.strip().lower() in real_answer_yes: + # The answer fact is known to be landmark + self.scores.append(True) + elif x.strip().lower() == "none": + # The answer is none, correct only if there are no known landmarks, + # since we only generate questions when that means that there are no non-trivial landmarks + self.scores.append(len(real_answer_yes) == 0) + else: + # All other cases the answer is incorrect + self.scores.append(False) + + return self.get_avg_score() + + +""" +Next Action task: generate an action that takes us closer to the goal. +answer: + (a) A list of applicable actions that are known to be correct answers + (b) A list of applicable actions that are known to be incorrect answers + (c) The rest of the applicable actions (maybe). +""" + + +class NextActionEvaluator(BaseEvaluator): + def get_score(self, ans, doc): + real_answer = doc["answer"] + real_answer_yes = [a.lower() for a in real_answer["yes"]] + real_answer_no = [a.lower() for a in real_answer["no"]] + real_answer_maybe = [a.lower() for a in real_answer["maybe"]] + # The cost of the optimal plan from the current state + opt = real_answer.get("opt", None) + for x in ans: + if x is None: + self.scores.append(False) + continue + action = x.strip().lower() + if action in real_answer_yes: + # Known to be correct + self.scores.append(True) + elif action in real_answer_no: + # Known to be incorrect + self.scores.append(False) + elif action not in real_answer_maybe: + # Not applicable, must be incorrect + self.scores.append(False) + else: + # Unknown, need to run a planner to check whether the state that results from applying the action is closer to the goal + # meaning has smaller optimal plan cost. + self.scores.append( + is_on_optimal_plan( + doc["PDDL_domain"].lower(), + doc["PDDL_problem"].lower(), + action, + opt, + ) + ) + + return self.get_avg_score() + + +""" +Progression task: generate the positive and negative effects of an action in the current state. +answer: + (a) A list of facts that were false and become true, when the action is applied + (b) A list of facts that were true and become false, when the action is applied +""" + + +class ProgressionEvaluator(BaseEvaluator): + def get_score(self, ans, doc): + real_answer = doc["answer"] + real_answer_pos = [a.lower() for a in real_answer["pos"]] + real_answer_neg = [a.lower() for a in real_answer["neg"]] + + for x in ans: + # The answer should be two lists. We allow for a single list and assume that the second one is empty (relaxed evaluation). + if x is None or len(x) > 2 or len(x) < 1: + self.scores.append(False) + else: + p = cleanup_answer(x[0]) + if len(x) == 2: + n = cleanup_answer(x[1]) + else: + # Assuming the last element is dropped because it is empty + n = [] + # Check if the answer is equal as sets to the correct answers. + ans = [set_equal(real_answer_pos, p), set_equal(real_answer_neg, n)] + self.scores.append(all(ans)) + + return self.get_avg_score() + + +""" +Reachability task: generate a valid fact that will never become true in any reachable state. +answer: A subset of facts that are known to be unreachable (not an exhaustive set). + It is empty only when we *know* that there are no such facts. +""" + + +class ReachabilityEvaluator(BaseEvaluator): + def get_score(self, ans, doc): + real_answer = doc["answer"] + real_answer = [f"({x.strip().lower()})" for x in real_answer] + + if len(real_answer) == 0: + # The correct answer is None + self.add_scores( + ["none" == x.strip().lower() if x is not None else False for x in ans] + ) + else: + for x in ans: + if x is None: + self.scores.append(False) + elif x.strip().lower() in real_answer: + # The answer is in the subset of stored correct answers + self.scores.append(True) + else: + # Need to run a planner on a task with the answer fact as the new goal + atom = x.strip().lower() + self.scores.append( + is_unsolvable_new_goal( + doc["PDDL_domain"].lower(), + doc["PDDL_problem"].lower(), + atom, + ) + ) + + return self.get_avg_score() + + +""" +Validation task: generate an index of the first inapplicable action in the given sequence. +answer: the correct index. +""" + + +class ValidationEvaluator(BaseEvaluator): + def get_score(self, ans, doc): + real_answer = str(doc["answer"]) + assert int(real_answer) >= 0, ( + f"The index must be non-negative, received {real_answer}" + ) + # Exact match + self.add_scores( + [ + real_answer.lower() == x.strip().lower() if x is not None else False + for x in ans + ] + ) + + return self.get_avg_score() + + +############################################################################## + + +def dump_item(item, **kwargs): + return json.dumps(item) + + +def parse_prediction(prediction): + try: + ans = json.loads(prediction.strip()) + response = ans.get("answer", None) + return response + except Exception as e: + print(f"Exception occurred {e}") + return prediction + + +@register_filter("ACP_grammar_filter") +class ACPGrammarFilter(RegexFilter): + """Filtering Index using""" + + def __init__(self, *args, **kwargs): + self.parser = ACPGrammarParser(kwargs["grammar_task"]) + self.clean = kwargs["clean"] if "clean" in kwargs else None + + def clean_pos_neg(self, resp): + # Check for Positive Effects and Negative Effects instead of separation + if check_prog_response(resp): + resp2 = resp.lower() + resp2 = resp2.replace("*", "") + resp2 = resp2.replace("positive effects", "[") + resp2 = resp2.replace("negative effects", "] [") + resp2 = resp2 + "]" + return resp2 + return resp + + def clean_simplified_plan(self, resp): + # Check for "simplified plan:" + if "simplified plan:" in resp.lower(): + resp2 = resp.lower() + resp2 = resp2.replace("*", "") + resp2 = resp2.split("simplified plan:")[1] + return resp2 + return resp + + def apply(self, resps, docs): + if self.clean == "pos_neg": + filtered_resps = [ + [self.parser.parse(self.clean_pos_neg(r)) for r in resp] + for resp in resps + ] + elif self.clean == "simplified plan": + filtered_resps = [ + [self.parser.parse(self.clean_simplified_plan(r)) for r in resp] + for resp in resps + ] + else: + filtered_resps = [[self.parser.parse(r) for r in resp] for resp in resps] + return filtered_resps + + +def process_acp_results(doc, results): + return {"score": get_evaluator(doc["group"]).get_score(results, doc)} + + +def get_score(references, predictions, **kwargs): + # print(f"References: {references}") + # print(f"Predictions: {predictions}") + data = json.loads(references[0].strip()) + real_ans = data["answer"] + task = data["group"] + + responses = [parse_prediction(prediction) for prediction in predictions] + + print(f"Real answer: {real_ans}") + print(f"Model answers: {responses}") + parser = ACPGrammarParser(get_grammar_task(task)) + ans = parse_ans(responses, parser, task) + + print(f"Parsed model answers: {ans}") + score = get_evaluator(task).get_score(ans, data) + + return {"get_score": score} diff --git a/lm-evaluation-harness/lm_eval/tasks/acpbench/gen_2shot_with_pddl/app.yaml b/lm-evaluation-harness/lm_eval/tasks/acpbench/gen_2shot_with_pddl/app.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c212924c04222f24918161e9fb46587705bd520d --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/acpbench/gen_2shot_with_pddl/app.yaml @@ -0,0 +1,23 @@ +task: acp_app_gen_with_pddl +dataset_name: acp_app_gen +include: _gen_yaml_2shot +fewshot_config: + sampler: first_n + samples: + - context: "A robot is in a grid and can only move to places that are connected to its current position. \nThe grid size is 5x5, and the locations are of the form fi-jf (e.g., f3-2f or f0-1f). The grid cells are connected to their neighbors (e.g., f1-2f is connected to the four neighbors f0-2f, f2-2f, f1-1f, and f1-3f). Some positions on the grid are locked and can be opened with a key of a matching shape. The robot has an arm that can pick up a key when the key is in same location as the robot and the arm is empty. \nThere are 2 keys in 0 different shapes: Key key0-0 is of shape shape0, Key key0-1 is of shape shape0. \nCurrently, the robot is at position f4-3f and its arm is empty. All the positions are open except the following: f2-0f has shape0 shaped lock, f4-2f has shape0 shaped lock. Key key0-0 is at position f3-1f. Key key0-1 is at position f1-3f. The available actions are: (unlock ?curpos ?lockpos ?key ?shape) - use the key ?key of shape ?shape to unlock the place ?lockpos from the current position ?curpos, (move ?curpos ?nextpos) - transition from the current position ?curpos to the next position ?nextpos, (pickup ?curpos ?key) - pick up key ?key from place ?curpos, (pickup-and-loose ?curpos ?newkey ?oldkey) - pick up the key ?newkey from the current position ?curpos and loose the key ?oldkey which is being held, and (putdown ?curpos ?key) - place the key ?key at the current position ?curpos." + question: "Generate the list of all ground actions that are applicable in this state." + answer: "[(move f4-3f f3-3f), (move f4-3f f4-4f)]" + PDDL_domain: "(define (domain grid)\n (:requirements :strips :typing)\n (:types key place shape - object)\n (:predicates (arm-empty) (at ?r - key ?x - place) (at-robot ?x - place) (conn ?x - place ?y - place) (holding ?k - key) (key-shape ?k - key ?s - shape) (lock-shape ?x - place ?s - shape) (locked ?x - place) (open ?x - place))\n (:action move\n :parameters (?curpos - place ?nextpos - place)\n :precondition (and (at-robot ?curpos) (conn ?curpos ?nextpos) (open ?nextpos))\n :effect (and (at-robot ?nextpos) (not (at-robot ?curpos)))\n )\n (:action pickup\n :parameters (?curpos - place ?key - key)\n :precondition (and (at-robot ?curpos) (at ?key ?curpos) (arm-empty))\n :effect (and (holding ?key) (not (at ?key ?curpos)) (not (arm-empty)))\n )\n (:action pickup-and-loose\n :parameters (?curpos - place ?newkey - key ?oldkey - key)\n :precondition (and (at-robot ?curpos) (holding ?oldkey) (at ?newkey ?curpos))\n :effect (and (holding ?newkey) (at ?oldkey ?curpos) (not (holding ?oldkey)) (not (at ?newkey ?curpos)))\n )\n (:action putdown\n :parameters (?curpos - place ?key - key)\n :precondition (and (at-robot ?curpos) (holding ?key))\n :effect (and (arm-empty) (at ?key ?curpos) (not (holding ?key)))\n )\n (:action unlock\n :parameters (?curpos - place ?lockpos - place ?key - key ?shape - shape)\n :precondition (and (conn ?curpos ?lockpos) (key-shape ?key ?shape) (lock-shape ?lockpos ?shape) (at-robot ?curpos) (locked ?lockpos) (holding ?key))\n :effect (and (open ?lockpos) (not (locked ?lockpos)))\n )\n)" + PDDL_problem: "(define (problem grid-x5-y5-t1-k2-l2-p100)\n (:domain grid)\n (:requirements :strips :typing)\n (:objects key0-0 key0-1 - key f0-0f f0-1f f0-2f f0-3f f0-4f f1-0f f1-1f f1-2f f1-3f f1-4f f2-0f f2-1f f2-2f f2-3f f2-4f f3-0f f3-1f f3-2f f3-3f f3-4f f4-0f f4-1f f4-2f f4-3f f4-4f - place shape0 - shape)\n (:init (arm-empty) (at key0-0 f3-1f) (at key0-1 f1-3f) (at-robot f4-3f) (conn f0-0f f0-1f) (conn f0-0f f1-0f) (conn f0-1f f0-0f) (conn f0-1f f0-2f) (conn f0-1f f1-1f) (conn f0-2f f0-1f) (conn f0-2f f0-3f) (conn f0-2f f1-2f) (conn f0-3f f0-2f) (conn f0-3f f0-4f) (conn f0-3f f1-3f) (conn f0-4f f0-3f) (conn f0-4f f1-4f) (conn f1-0f f0-0f) (conn f1-0f f1-1f) (conn f1-0f f2-0f) (conn f1-1f f0-1f) (conn f1-1f f1-0f) (conn f1-1f f1-2f) (conn f1-1f f2-1f) (conn f1-2f f0-2f) (conn f1-2f f1-1f) (conn f1-2f f1-3f) (conn f1-2f f2-2f) (conn f1-3f f0-3f) (conn f1-3f f1-2f) (conn f1-3f f1-4f) (conn f1-3f f2-3f) (conn f1-4f f0-4f) (conn f1-4f f1-3f) (conn f1-4f f2-4f) (conn f2-0f f1-0f) (conn f2-0f f2-1f) (conn f2-0f f3-0f) (conn f2-1f f1-1f) (conn f2-1f f2-0f) (conn f2-1f f2-2f) (conn f2-1f f3-1f) (conn f2-2f f1-2f) (conn f2-2f f2-1f) (conn f2-2f f2-3f) (conn f2-2f f3-2f) (conn f2-3f f1-3f) (conn f2-3f f2-2f) (conn f2-3f f2-4f) (conn f2-3f f3-3f) (conn f2-4f f1-4f) (conn f2-4f f2-3f) (conn f2-4f f3-4f) (conn f3-0f f2-0f) (conn f3-0f f3-1f) (conn f3-0f f4-0f) (conn f3-1f f2-1f) (conn f3-1f f3-0f) (conn f3-1f f3-2f) (conn f3-1f f4-1f) (conn f3-2f f2-2f) (conn f3-2f f3-1f) (conn f3-2f f3-3f) (conn f3-2f f4-2f) (conn f3-3f f2-3f) (conn f3-3f f3-2f) (conn f3-3f f3-4f) (conn f3-3f f4-3f) (conn f3-4f f2-4f) (conn f3-4f f3-3f) (conn f3-4f f4-4f) (conn f4-0f f3-0f) (conn f4-0f f4-1f) (conn f4-1f f3-1f) (conn f4-1f f4-0f) (conn f4-1f f4-2f) (conn f4-2f f3-2f) (conn f4-2f f4-1f) (conn f4-2f f4-3f) (conn f4-3f f3-3f) (conn f4-3f f4-2f) (conn f4-3f f4-4f) (conn f4-4f f3-4f) (conn f4-4f f4-3f) (key-shape key0-0 shape0) (key-shape key0-1 shape0) (lock-shape f2-0f shape0) (lock-shape f4-2f shape0) (locked f2-0f) (locked f4-2f) (open f0-0f) (open f0-1f) (open f0-2f) (open f0-3f) (open f0-4f) (open f1-0f) (open f1-1f) (open f1-2f) (open f1-3f) (open f1-4f) (open f2-1f) (open f2-2f) (open f2-3f) (open f2-4f) (open f3-0f) (open f3-1f) (open f3-2f) (open f3-3f) (open f3-4f) (open f4-0f) (open f4-1f) (open f4-3f) (open f4-4f))\n (:goal (and (at key0-0 f2-0f) (at key0-1 f1-3f)))\n)" + - context: "There are several cities, each containing several locations, some of which are airports. There are also trucks, which can drive within a single city, and airplanes, which can fly between airports. The goal is to get some packages from various locations to various new locations. \nThere are 2 trucks and 1 airplane, as well as 4 packages. There are 4 locations across 2 cities. \nThe locations are in cities as follows: l0-1 and l0-0 are in c0; l1-1 and l1-0 are in c1. \nCurrently, t1 is at l1-0, p0, a0, t0, and p3 are at l0-0, p1 and p2 are in t1. The available actions are: (load-truck ?obj ?truck ?loc) - load the object ?obj from location ?loc into the truck ?truck, (load-airplane ?obj ?airplane ?loc) - place the object ?obj onto the airplane ?airplane at location ?loc, (unload-truck ?obj ?truck ?loc) - remove the object ?obj from the truck ?truck and place it on the location ?loc, (unload-airplane ?obj ?airplane ?loc) - unload object ?obj from airplane ?airplane at location ?loc, (drive-truck ?truck ?loc-from ?loc-to ?city) - navigate the truck ?truck from location ?loc-from in city ?city to location ?loc-to in the same city, and (fly-airplane ?airplane ?loc-from ?loc-to) - fly airplane ?airplane from airport ?loc-from to airport ?loc-to." + question: "Generate the list of all ground actions that are applicable in this state." + answer: "[(unload-truck p2 t1 l1-0), (drive-truck t0 l0-0 l0-0 c0), (load-airplane p0 a0 l0-0), (load-truck p0 t0 l0-0), (unload-truck p1 t1 l1-0), (drive-truck t1 l1-0 l1-0 c1), (drive-truck t0 l0-0 l0-1 c0), (drive-truck t1 l1-0 l1-1 c1), (fly-airplane a0 l0-0 l0-0), (load-truck p3 t0 l0-0), (fly-airplane a0 l0-0 l1-0), (load-airplane p3 a0 l0-0)]" + PDDL_domain: "(define (domain logistics-strips)\n (:requirements :strips :typing) \n\n (:types \n location locatable city - object \n package movable - locatable\n airport - location\n airplane truck - movable \n )\t\t\n \n (:predicates \t\n\t\t(at ?obj - locatable ?loc - location)\n\t\t(in ?obj1 - package ?obj2 - movable)\n\t\t(in-city ?obj - location ?city - city))\n\n\n(:action LOAD-TRUCK\n :parameters\n (?obj - package\n ?truck - truck\n ?loc - location)\n :precondition\n (and \n (at ?truck ?loc) (at ?obj ?loc))\n :effect\n (and (not (at ?obj ?loc)) (in ?obj ?truck)))\n\n(:action LOAD-AIRPLANE\n :parameters\n (?obj - package\n ?airplane - airplane\n ?loc - location)\n :precondition\n (and \n (at ?obj ?loc) (at ?airplane ?loc))\n :effect\n (and (not (at ?obj ?loc)) (in ?obj ?airplane)))\n\n\n\n(:action UNLOAD-TRUCK\n :parameters\n (?obj - package\n ?truck - truck\n ?loc - location)\n :precondition\n (and \n (at ?truck ?loc) (in ?obj ?truck))\n :effect\n (and (not (in ?obj ?truck)) (at ?obj ?loc)))\n\n(:action UNLOAD-AIRPLANE\n :parameters\n (?obj - package\n ?airplane - airplane\n ?loc - location)\n :precondition\n (and \n (in ?obj ?airplane) (at ?airplane ?loc))\n :effect\n (and (not (in ?obj ?airplane)) (at ?obj ?loc)))\n\n(:action DRIVE-TRUCK\n :parameters\n (?truck - truck\n ?loc-from - location\n ?loc-to - location\n ?city - city)\n :precondition\n (and \n (at ?truck ?loc-from)\n (in-city ?loc-from ?city)\n (in-city ?loc-to ?city))\n :effect\n (and (not (at ?truck ?loc-from)) (at ?truck ?loc-to)))\n\n(:action FLY-AIRPLANE\n :parameters\n (?airplane - airplane\n ?loc-from - airport\n ?loc-to - airport)\n :precondition\n (and \n\t(at ?airplane ?loc-from))\n :effect\n (and (not (at ?airplane ?loc-from)) (at ?airplane ?loc-to)))\n)" + PDDL_problem: "(define (problem logistics-c2-s2-p4-a1)\n (:domain logistics-strips)\n (:requirements :strips :typing)\n (:objects a0 - airplane l0-0 l1-0 - airport c0 c1 - city l0-1 l1-1 - location p0 p1 p2 p3 - package t0 t1 - truck)\n (:init (at a0 l0-0) (at p0 l0-0) (at p3 l0-0) (at t0 l0-0) (at t1 l1-0) (in p1 t1) (in p2 t1) (in-city l0-0 c0) (in-city l0-1 c0) (in-city l1-0 c1) (in-city l1-1 c1))\n (:goal (and (at p0 l0-0) (at p1 l1-0) (at p2 l1-0) (at p3 l0-1)))\n)" +doc_to_text: "# PDDL DOMAIN \n\n```\n{{PDDL_domain}}\n```\n\n# PDDL PROBLEM \n\n```\n{{PDDL_problem}}\n```\n\n**Question**: {{context}} {{question}} Each action starts with an opening parenthesis and ends with closing parenthesis. Provide only the actions. \n**Final Answer**:" +filter_list: + - name: "acp_grammar_parse" + filter: + - function: "ACP_grammar_filter" + grammar_task: "action_list" + - function: "take_first" diff --git a/lm-evaluation-harness/lm_eval/tasks/acpbench/gen_2shot_with_pddl/just.yaml b/lm-evaluation-harness/lm_eval/tasks/acpbench/gen_2shot_with_pddl/just.yaml new file mode 100644 index 0000000000000000000000000000000000000000..9685b8b8f3c3c9274ca83f9daacff26de7d26200 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/acpbench/gen_2shot_with_pddl/just.yaml @@ -0,0 +1,24 @@ +task: acp_just_gen_with_pddl +dataset_name: acp_just_gen +include: _gen_yaml_2shot +fewshot_config: + sampler: first_n + samples: + - context: "A robot is in a grid and can only move to places that are connected to its current position. \nThe grid size is 5x5, and the locations are of the form fi-jf (e.g., f3-2f or f0-1f). The grid cells are connected to their neighbors (e.g., f1-2f is connected to the four neighbors f0-2f, f2-2f, f1-1f, and f1-3f). Some positions on the grid are locked and can be opened with a key of a matching shape. The robot has an arm that can pick up a key when the key is in same location as the robot and the arm is empty. \nThere are 2 keys in 0 different shapes: Key key0-0 is of shape shape0, Key key0-1 is of shape shape0. \nCurrently, the robot is at position f3-3f and its arm is empty. All the positions are open except the following: f4-2f has shape0 shaped lock, f2-0f has shape0 shaped lock. Key key0-1 is at position f1-3f. Key key0-0 is at position f2-2f. The available actions are: (unlock ?curpos ?lockpos ?key ?shape) - unlock place ?lockpos with key ?key of shape ?shape from current position place ?curpos, (move ?curpos ?nextpos) - transition from the current position ?curpos to the next position ?nextpos, (pickup ?curpos ?key) - acquire the key ?key from the place ?curpos, (pickup-and-loose ?curpos ?newkey ?oldkey) - pick up the key ?newkey from the current position ?curpos and loose the key ?oldkey which is being held, and (putdown ?curpos ?key) - place the key ?key at the current position place ?curpos. The goal is to reach a state where the following facts hold: Key key0-1 is at f1-3f location and Key key0-0 is at f2-0f location." + question: "Simplify the plan \"(move f3-3f f3-2f) (move f3-2f f2-2f) (pickup f2-2f key0-0) (move f2-2f f2-1f) (putdown f2-1f key0-0) (pickup f2-1f key0-0) (unlock f2-1f f2-0f key0-0 shape0) (move f2-1f f2-0f) (putdown f2-0f key0-0)\" by removing either a single action or a pair of consecutive actions, while still maintaining a valid plan. Provide the resulting simplified plan." + answer: "[(move f3-3f f3-2f), (move f3-2f f2-2f), (pickup f2-2f key0-0), (move f2-2f f2-1f), (unlock f2-1f f2-0f key0-0 shape0), (move f2-1f f2-0f), (putdown f2-0f key0-0)]" + PDDL_domain: "(define (domain grid)\n (:requirements :strips :typing)\n (:types key place shape - object)\n (:predicates (arm-empty) (at ?r - key ?x - place) (at-robot ?x - place) (conn ?x - place ?y - place) (holding ?k - key) (key-shape ?k - key ?s - shape) (lock-shape ?x - place ?s - shape) (locked ?x - place) (open ?x - place))\n (:action move\n :parameters (?curpos - place ?nextpos - place)\n :precondition (and (at-robot ?curpos) (conn ?curpos ?nextpos) (open ?nextpos))\n :effect (and (at-robot ?nextpos) (not (at-robot ?curpos)))\n )\n (:action pickup\n :parameters (?curpos - place ?key - key)\n :precondition (and (at-robot ?curpos) (at ?key ?curpos) (arm-empty))\n :effect (and (holding ?key) (not (at ?key ?curpos)) (not (arm-empty)))\n )\n (:action pickup-and-loose\n :parameters (?curpos - place ?newkey - key ?oldkey - key)\n :precondition (and (at-robot ?curpos) (holding ?oldkey) (at ?newkey ?curpos))\n :effect (and (holding ?newkey) (at ?oldkey ?curpos) (not (holding ?oldkey)) (not (at ?newkey ?curpos)))\n )\n (:action putdown\n :parameters (?curpos - place ?key - key)\n :precondition (and (at-robot ?curpos) (holding ?key))\n :effect (and (arm-empty) (at ?key ?curpos) (not (holding ?key)))\n )\n (:action unlock\n :parameters (?curpos - place ?lockpos - place ?key - key ?shape - shape)\n :precondition (and (conn ?curpos ?lockpos) (key-shape ?key ?shape) (lock-shape ?lockpos ?shape) (at-robot ?curpos) (locked ?lockpos) (holding ?key))\n :effect (and (open ?lockpos) (not (locked ?lockpos)))\n )\n)" + PDDL_problem: "(define (problem grid-x5-y5-t1-k2-l2-p100)\n (:domain grid)\n (:requirements :strips :typing)\n (:objects key0-0 key0-1 - key f0-0f f0-1f f0-2f f0-3f f0-4f f1-0f f1-1f f1-2f f1-3f f1-4f f2-0f f2-1f f2-2f f2-3f f2-4f f3-0f f3-1f f3-2f f3-3f f3-4f f4-0f f4-1f f4-2f f4-3f f4-4f - place shape0 - shape)\n (:init (arm-empty) (at key0-0 f2-2f) (at key0-1 f1-3f) (at-robot f3-3f) (conn f0-0f f0-1f) (conn f0-0f f1-0f) (conn f0-1f f0-0f) (conn f0-1f f0-2f) (conn f0-1f f1-1f) (conn f0-2f f0-1f) (conn f0-2f f0-3f) (conn f0-2f f1-2f) (conn f0-3f f0-2f) (conn f0-3f f0-4f) (conn f0-3f f1-3f) (conn f0-4f f0-3f) (conn f0-4f f1-4f) (conn f1-0f f0-0f) (conn f1-0f f1-1f) (conn f1-0f f2-0f) (conn f1-1f f0-1f) (conn f1-1f f1-0f) (conn f1-1f f1-2f) (conn f1-1f f2-1f) (conn f1-2f f0-2f) (conn f1-2f f1-1f) (conn f1-2f f1-3f) (conn f1-2f f2-2f) (conn f1-3f f0-3f) (conn f1-3f f1-2f) (conn f1-3f f1-4f) (conn f1-3f f2-3f) (conn f1-4f f0-4f) (conn f1-4f f1-3f) (conn f1-4f f2-4f) (conn f2-0f f1-0f) (conn f2-0f f2-1f) (conn f2-0f f3-0f) (conn f2-1f f1-1f) (conn f2-1f f2-0f) (conn f2-1f f2-2f) (conn f2-1f f3-1f) (conn f2-2f f1-2f) (conn f2-2f f2-1f) (conn f2-2f f2-3f) (conn f2-2f f3-2f) (conn f2-3f f1-3f) (conn f2-3f f2-2f) (conn f2-3f f2-4f) (conn f2-3f f3-3f) (conn f2-4f f1-4f) (conn f2-4f f2-3f) (conn f2-4f f3-4f) (conn f3-0f f2-0f) (conn f3-0f f3-1f) (conn f3-0f f4-0f) (conn f3-1f f2-1f) (conn f3-1f f3-0f) (conn f3-1f f3-2f) (conn f3-1f f4-1f) (conn f3-2f f2-2f) (conn f3-2f f3-1f) (conn f3-2f f3-3f) (conn f3-2f f4-2f) (conn f3-3f f2-3f) (conn f3-3f f3-2f) (conn f3-3f f3-4f) (conn f3-3f f4-3f) (conn f3-4f f2-4f) (conn f3-4f f3-3f) (conn f3-4f f4-4f) (conn f4-0f f3-0f) (conn f4-0f f4-1f) (conn f4-1f f3-1f) (conn f4-1f f4-0f) (conn f4-1f f4-2f) (conn f4-2f f3-2f) (conn f4-2f f4-1f) (conn f4-2f f4-3f) (conn f4-3f f3-3f) (conn f4-3f f4-2f) (conn f4-3f f4-4f) (conn f4-4f f3-4f) (conn f4-4f f4-3f) (key-shape key0-0 shape0) (key-shape key0-1 shape0) (lock-shape f2-0f shape0) (lock-shape f4-2f shape0) (locked f2-0f) (locked f4-2f) (open f0-0f) (open f0-1f) (open f0-2f) (open f0-3f) (open f0-4f) (open f1-0f) (open f1-1f) (open f1-2f) (open f1-3f) (open f1-4f) (open f2-1f) (open f2-2f) (open f2-3f) (open f2-4f) (open f3-0f) (open f3-1f) (open f3-2f) (open f3-3f) (open f3-4f) (open f4-0f) (open f4-1f) (open f4-3f) (open f4-4f))\n (:goal (and (at key0-0 f2-0f) (at key0-1 f1-3f)))\n)" + - context: "There are several cities, each containing several locations, some of which are airports. There are also trucks, which can drive within a single city, and airplanes, which can fly between airports. The goal is to get some packages from various locations to various new locations. \nThere are 2 trucks and 1 airplane, as well as 4 packages. There are 4 locations across 2 cities. \nThe locations are in cities as follows: l1-0 and l1-1 are in c1; l0-1 and l0-0 are in c0. \nCurrently, p3, p2, and p1 are at l1-0, t0 is at l0-1, p0 and t1 are at l1-1, a0 is at l0-0. The available actions are: (load-truck ?obj ?truck ?loc) - place the object ?obj into the truck ?truck at location ?loc, (load-airplane ?obj ?airplane ?loc) - load the object ?obj from location ?loc into the airplane ?airplane, (unload-truck ?obj ?truck ?loc) - unload the object ?obj from the truck ?truck at location ?loc, (unload-airplane ?obj ?airplane ?loc) - remove the object ?obj from the airplane ?airplane and place it on the location ?loc, (drive-truck ?truck ?loc-from ?loc-to ?city) - navigate the truck ?truck which is in location ?loc-from in city ?city to another location ?loc-to in the same city, and (fly-airplane ?airplane ?loc-from ?loc-to) - fly the airplane ?airplane from airport ?loc-from to airport ?loc-to. The goal is to reach a state where the following facts hold: p2 is at l1-0, p0 is at l0-0, p3 is at l0-1, and p1 is at l1-0." + question: "Simplify the plan \"(load-truck p0 t1 l1-1) (unload-truck p0 t1 l1-1) (load-truck p0 t1 l1-1) (drive-truck t1 l1-1 l1-0 c1) (unload-truck p0 t1 l1-0) (fly-airplane a0 l0-0 l1-0) (load-airplane p0 a0 l1-0) (load-airplane p3 a0 l1-0) (fly-airplane a0 l1-0 l0-0) (unload-airplane p0 a0 l0-0) (unload-airplane p3 a0 l0-0) (drive-truck t0 l0-1 l0-0 c0) (load-truck p3 t0 l0-0) (drive-truck t0 l0-0 l0-1 c0) (unload-truck p3 t0 l0-1)\" by removing either a single action or a pair of consecutive actions, while still maintaining a valid plan. Provide the resulting simplified plan." + answer: "[(load-truck p0 t1 l1-1), (drive-truck t1 l1-1 l1-0 c1), (unload-truck p0 t1 l1-0), (fly-airplane a0 l0-0 l1-0), (load-airplane p0 a0 l1-0), (load-airplane p3 a0 l1-0), (fly-airplane a0 l1-0 l0-0), (unload-airplane p0 a0 l0-0), (unload-airplane p3 a0 l0-0), (drive-truck t0 l0-1 l0-0 c0), (load-truck p3 t0 l0-0), (drive-truck t0 l0-0 l0-1 c0), (unload-truck p3 t0 l0-1)]" + PDDL_domain: "(define (domain logistics-strips)\n (:requirements :strips :typing) \n\n (:types \n location locatable city - object \n package movable - locatable\n airport - location\n airplane truck - movable \n )\t\t\n \n (:predicates \t\n\t\t(at ?obj - locatable ?loc - location)\n\t\t(in ?obj1 - package ?obj2 - movable)\n\t\t(in-city ?obj - location ?city - city))\n\n\n(:action LOAD-TRUCK\n :parameters\n (?obj - package\n ?truck - truck\n ?loc - location)\n :precondition\n (and \n (at ?truck ?loc) (at ?obj ?loc))\n :effect\n (and (not (at ?obj ?loc)) (in ?obj ?truck)))\n\n(:action LOAD-AIRPLANE\n :parameters\n (?obj - package\n ?airplane - airplane\n ?loc - location)\n :precondition\n (and \n (at ?obj ?loc) (at ?airplane ?loc))\n :effect\n (and (not (at ?obj ?loc)) (in ?obj ?airplane)))\n\n\n\n(:action UNLOAD-TRUCK\n :parameters\n (?obj - package\n ?truck - truck\n ?loc - location)\n :precondition\n (and \n (at ?truck ?loc) (in ?obj ?truck))\n :effect\n (and (not (in ?obj ?truck)) (at ?obj ?loc)))\n\n(:action UNLOAD-AIRPLANE\n :parameters\n (?obj - package\n ?airplane - airplane\n ?loc - location)\n :precondition\n (and \n (in ?obj ?airplane) (at ?airplane ?loc))\n :effect\n (and (not (in ?obj ?airplane)) (at ?obj ?loc)))\n\n(:action DRIVE-TRUCK\n :parameters\n (?truck - truck\n ?loc-from - location\n ?loc-to - location\n ?city - city)\n :precondition\n (and \n (at ?truck ?loc-from)\n (in-city ?loc-from ?city)\n (in-city ?loc-to ?city))\n :effect\n (and (not (at ?truck ?loc-from)) (at ?truck ?loc-to)))\n\n(:action FLY-AIRPLANE\n :parameters\n (?airplane - airplane\n ?loc-from - airport\n ?loc-to - airport)\n :precondition\n (and \n\t(at ?airplane ?loc-from))\n :effect\n (and (not (at ?airplane ?loc-from)) (at ?airplane ?loc-to)))\n)" + PDDL_problem: "(define (problem logistics-c2-s2-p4-a1)\n (:domain logistics-strips)\n (:requirements :strips :typing)\n (:objects a0 - airplane l0-0 l1-0 - airport c0 c1 - city l0-1 l1-1 - location p0 p1 p2 p3 - package t0 t1 - truck)\n (:init (at a0 l0-0) (at p0 l1-1) (at p1 l1-0) (at p2 l1-0) (at p3 l1-0) (at t0 l0-1) (at t1 l1-1) (in-city l0-0 c0) (in-city l0-1 c0) (in-city l1-0 c1) (in-city l1-1 c1))\n (:goal (and (at p0 l0-0) (at p1 l1-0) (at p2 l1-0) (at p3 l0-1)))\n)" +doc_to_text: "# PDDL DOMAIN \n\n```\n{{PDDL_domain}}\n```\n\n# PDDL PROBLEM \n\n```\n{{PDDL_problem}}\n```\n\n**Question**: {{context}} {{question}} **Final Answer**:" +filter_list: + - name: "acp_grammar_parse" + filter: + - function: "ACP_grammar_filter" + grammar_task: "action_list" + clean: "simplified plan" + - function: "take_first" diff --git a/lm-evaluation-harness/lm_eval/tasks/acpbench/gen_2shot_with_pddl/next_act.yaml b/lm-evaluation-harness/lm_eval/tasks/acpbench/gen_2shot_with_pddl/next_act.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0f43ca61c73417762d250c83d64e1d7b1bfde37f --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/acpbench/gen_2shot_with_pddl/next_act.yaml @@ -0,0 +1,23 @@ +task: acp_nexta_gen_with_pddl +dataset_name: acp_nexta_gen +include: _gen_yaml_2shot +fewshot_config: + sampler: first_n + samples: + - context: "A robot is in a grid and can only move to places that are connected to its current position. \nThe grid size is 5x5, and the locations are of the form fi-jf (e.g., f3-2f or f0-1f). The grid cells are connected to their neighbors (e.g., f1-2f is connected to the four neighbors f0-2f, f2-2f, f1-1f, and f1-3f). Some positions on the grid are locked and can be opened with a key of a matching shape. The robot has an arm that can pick up a key when the key is in same location as the robot and the arm is empty. \nThere are 2 keys in 0 different shapes: Key key0-0 is of shape shape0, Key key0-1 is of shape shape0. \nCurrently, the robot is at position f1-1f and its arm is empty. All the positions are open except the following: f4-2f has shape0 shaped lock. Key key0-0 is at position f1-0f. Key key0-1 is at position f1-3f. The goal is to reach a state where the following facts hold: Key key0-1 is at f1-3f location and Key key0-0 is at f2-0f location. The available actions are: (unlock ?curpos ?lockpos ?key ?shape) - use the key ?key of shape ?shape to unlock the place ?lockpos from the current position ?curpos, (move ?curpos ?nextpos) - move to place ?nextpos from place ?curpos, (pickup ?curpos ?key) - retrieve the key ?key from its current position ?curpos, (pickup-and-loose ?curpos ?newkey ?oldkey) - pick up key ?newkey at current position place ?curpos and loose key ?oldkey being held, and (putdown ?curpos ?key) - put the key ?key at the current position place ?curpos." + question: "What is the next action that takes us towards the goal?" + answer: "(move f1-1f f1-0f)" + PDDL_domain: "(define (domain grid)\n (:requirements :strips :typing)\n (:types key place shape - object)\n (:predicates (arm-empty) (at ?r - key ?x - place) (at-robot ?x - place) (conn ?x - place ?y - place) (holding ?k - key) (key-shape ?k - key ?s - shape) (lock-shape ?x - place ?s - shape) (locked ?x - place) (open ?x - place))\n (:action move\n :parameters (?curpos - place ?nextpos - place)\n :precondition (and (at-robot ?curpos) (conn ?curpos ?nextpos) (open ?nextpos))\n :effect (and (at-robot ?nextpos) (not (at-robot ?curpos)))\n )\n (:action pickup\n :parameters (?curpos - place ?key - key)\n :precondition (and (at-robot ?curpos) (at ?key ?curpos) (arm-empty))\n :effect (and (holding ?key) (not (at ?key ?curpos)) (not (arm-empty)))\n )\n (:action pickup-and-loose\n :parameters (?curpos - place ?newkey - key ?oldkey - key)\n :precondition (and (at-robot ?curpos) (holding ?oldkey) (at ?newkey ?curpos))\n :effect (and (holding ?newkey) (at ?oldkey ?curpos) (not (holding ?oldkey)) (not (at ?newkey ?curpos)))\n )\n (:action putdown\n :parameters (?curpos - place ?key - key)\n :precondition (and (at-robot ?curpos) (holding ?key))\n :effect (and (arm-empty) (at ?key ?curpos) (not (holding ?key)))\n )\n (:action unlock\n :parameters (?curpos - place ?lockpos - place ?key - key ?shape - shape)\n :precondition (and (conn ?curpos ?lockpos) (key-shape ?key ?shape) (lock-shape ?lockpos ?shape) (at-robot ?curpos) (locked ?lockpos) (holding ?key))\n :effect (and (open ?lockpos) (not (locked ?lockpos)))\n )\n)" + PDDL_problem: "(define (problem grid-x5-y5-t1-k2-l2-p100)\n (:domain grid)\n (:requirements :strips :typing)\n (:objects key0-0 key0-1 - key f0-0f f0-1f f0-2f f0-3f f0-4f f1-0f f1-1f f1-2f f1-3f f1-4f f2-0f f2-1f f2-2f f2-3f f2-4f f3-0f f3-1f f3-2f f3-3f f3-4f f4-0f f4-1f f4-2f f4-3f f4-4f - place shape0 - shape)\n (:init (arm-empty) (at key0-0 f1-0f) (at key0-1 f1-3f) (at-robot f1-1f) (conn f0-0f f0-1f) (conn f0-0f f1-0f) (conn f0-1f f0-0f) (conn f0-1f f0-2f) (conn f0-1f f1-1f) (conn f0-2f f0-1f) (conn f0-2f f0-3f) (conn f0-2f f1-2f) (conn f0-3f f0-2f) (conn f0-3f f0-4f) (conn f0-3f f1-3f) (conn f0-4f f0-3f) (conn f0-4f f1-4f) (conn f1-0f f0-0f) (conn f1-0f f1-1f) (conn f1-0f f2-0f) (conn f1-1f f0-1f) (conn f1-1f f1-0f) (conn f1-1f f1-2f) (conn f1-1f f2-1f) (conn f1-2f f0-2f) (conn f1-2f f1-1f) (conn f1-2f f1-3f) (conn f1-2f f2-2f) (conn f1-3f f0-3f) (conn f1-3f f1-2f) (conn f1-3f f1-4f) (conn f1-3f f2-3f) (conn f1-4f f0-4f) (conn f1-4f f1-3f) (conn f1-4f f2-4f) (conn f2-0f f1-0f) (conn f2-0f f2-1f) (conn f2-0f f3-0f) (conn f2-1f f1-1f) (conn f2-1f f2-0f) (conn f2-1f f2-2f) (conn f2-1f f3-1f) (conn f2-2f f1-2f) (conn f2-2f f2-1f) (conn f2-2f f2-3f) (conn f2-2f f3-2f) (conn f2-3f f1-3f) (conn f2-3f f2-2f) (conn f2-3f f2-4f) (conn f2-3f f3-3f) (conn f2-4f f1-4f) (conn f2-4f f2-3f) (conn f2-4f f3-4f) (conn f3-0f f2-0f) (conn f3-0f f3-1f) (conn f3-0f f4-0f) (conn f3-1f f2-1f) (conn f3-1f f3-0f) (conn f3-1f f3-2f) (conn f3-1f f4-1f) (conn f3-2f f2-2f) (conn f3-2f f3-1f) (conn f3-2f f3-3f) (conn f3-2f f4-2f) (conn f3-3f f2-3f) (conn f3-3f f3-2f) (conn f3-3f f3-4f) (conn f3-3f f4-3f) (conn f3-4f f2-4f) (conn f3-4f f3-3f) (conn f3-4f f4-4f) (conn f4-0f f3-0f) (conn f4-0f f4-1f) (conn f4-1f f3-1f) (conn f4-1f f4-0f) (conn f4-1f f4-2f) (conn f4-2f f3-2f) (conn f4-2f f4-1f) (conn f4-2f f4-3f) (conn f4-3f f3-3f) (conn f4-3f f4-2f) (conn f4-3f f4-4f) (conn f4-4f f3-4f) (conn f4-4f f4-3f) (key-shape key0-0 shape0) (key-shape key0-1 shape0) (lock-shape f2-0f shape0) (lock-shape f4-2f shape0) (locked f4-2f) (open f0-0f) (open f0-1f) (open f0-2f) (open f0-3f) (open f0-4f) (open f1-0f) (open f1-1f) (open f1-2f) (open f1-3f) (open f1-4f) (open f2-0f) (open f2-1f) (open f2-2f) (open f2-3f) (open f2-4f) (open f3-0f) (open f3-1f) (open f3-2f) (open f3-3f) (open f3-4f) (open f4-0f) (open f4-1f) (open f4-3f) (open f4-4f))\n (:goal (and (at key0-0 f2-0f) (at key0-1 f1-3f)))\n)" + - context: "There are several cities, each containing several locations, some of which are airports. There are also trucks, which can drive within a single city, and airplanes, which can fly between airports. The goal is to get some packages from various locations to various new locations. \nThere are 2 trucks and 1 airplane, as well as 4 packages. There are 4 locations across 2 cities. \nThe locations are in cities as follows: l1-1 and l1-0 are in c1; l0-0 and l0-1 are in c0. \nCurrently, p1, p3, t1, p2, and a0 are at l1-0, t0 is at l0-0, p0 is in a0. The goal is to reach a state where the following facts hold: p1 is at l1-0, p3 is at l0-1, p0 is at l0-0, and p2 is at l1-0. The available actions are: (load-truck ?obj ?truck ?loc) - load the object ?obj from location ?loc into the truck ?truck, (load-airplane ?obj ?airplane ?loc) - load the object ?obj from location ?loc onto the airplane ?airplane, (unload-truck ?obj ?truck ?loc) - unload the object ?obj from the truck ?truck at location ?loc, (unload-airplane ?obj ?airplane ?loc) - offload the object ?obj from the airplane ?airplane at location ?loc, (drive-truck ?truck ?loc-from ?loc-to ?city) - navigate the truck ?truck which is in location ?loc-from in city ?city to another location ?loc-to in the same city, and (fly-airplane ?airplane ?loc-from ?loc-to) - fly the airplane ?airplane from airport ?loc-from to airport ?loc-to." + question: "What is the next action that takes us towards the goal?" + answer: "(load-airplane p3 a0 l1-0)" + PDDL_domain: "(define (domain logistics-strips)\n (:requirements :strips :typing) \n\n (:types \n location locatable city - object \n package movable - locatable\n airport - location\n airplane truck - movable \n )\t\t\n \n (:predicates \t\n\t\t(at ?obj - locatable ?loc - location)\n\t\t(in ?obj1 - package ?obj2 - movable)\n\t\t(in-city ?obj - location ?city - city))\n\n\n(:action LOAD-TRUCK\n :parameters\n (?obj - package\n ?truck - truck\n ?loc - location)\n :precondition\n (and \n (at ?truck ?loc) (at ?obj ?loc))\n :effect\n (and (not (at ?obj ?loc)) (in ?obj ?truck)))\n\n(:action LOAD-AIRPLANE\n :parameters\n (?obj - package\n ?airplane - airplane\n ?loc - location)\n :precondition\n (and \n (at ?obj ?loc) (at ?airplane ?loc))\n :effect\n (and (not (at ?obj ?loc)) (in ?obj ?airplane)))\n\n\n\n(:action UNLOAD-TRUCK\n :parameters\n (?obj - package\n ?truck - truck\n ?loc - location)\n :precondition\n (and \n (at ?truck ?loc) (in ?obj ?truck))\n :effect\n (and (not (in ?obj ?truck)) (at ?obj ?loc)))\n\n(:action UNLOAD-AIRPLANE\n :parameters\n (?obj - package\n ?airplane - airplane\n ?loc - location)\n :precondition\n (and \n (in ?obj ?airplane) (at ?airplane ?loc))\n :effect\n (and (not (in ?obj ?airplane)) (at ?obj ?loc)))\n\n(:action DRIVE-TRUCK\n :parameters\n (?truck - truck\n ?loc-from - location\n ?loc-to - location\n ?city - city)\n :precondition\n (and \n (at ?truck ?loc-from)\n (in-city ?loc-from ?city)\n (in-city ?loc-to ?city))\n :effect\n (and (not (at ?truck ?loc-from)) (at ?truck ?loc-to)))\n\n(:action FLY-AIRPLANE\n :parameters\n (?airplane - airplane\n ?loc-from - airport\n ?loc-to - airport)\n :precondition\n (and \n\t(at ?airplane ?loc-from))\n :effect\n (and (not (at ?airplane ?loc-from)) (at ?airplane ?loc-to)))\n)" + PDDL_problem: "(define (problem logistics-c2-s2-p4-a1)\n (:domain logistics-strips)\n (:requirements :strips :typing)\n (:objects a0 - airplane l0-0 l1-0 - airport c0 c1 - city l0-1 l1-1 - location p0 p1 p2 p3 - package t0 t1 - truck)\n (:init (at a0 l1-0) (at p1 l1-0) (at p2 l1-0) (at p3 l1-0) (at t0 l0-0) (at t1 l1-0) (in p0 a0) (in-city l0-0 c0) (in-city l0-1 c0) (in-city l1-0 c1) (in-city l1-1 c1))\n (:goal (and (at p0 l0-0) (at p1 l1-0) (at p2 l1-0) (at p3 l0-1)))\n)" +doc_to_text: "# PDDL DOMAIN \n\n```\n{{PDDL_domain}}\n```\n\n# PDDL PROBLEM \n\n```\n{{PDDL_problem}}\n```\n\n**Question**: {{context}} {{question}} Each action starts with an opening parenthesis and ends with closing parenthesis. Provide only the action. **Final Answer**:" +filter_list: + - name: "acp_grammar_parse" + filter: + - function: "ACP_grammar_filter" + grammar_task: "action_name" + - function: "take_first" diff --git a/lm-evaluation-harness/lm_eval/tasks/acpbench/gen_2shot_with_pddl/prog.yaml b/lm-evaluation-harness/lm_eval/tasks/acpbench/gen_2shot_with_pddl/prog.yaml new file mode 100644 index 0000000000000000000000000000000000000000..545c56ee2bef5fec008217e434385f2b0c8f1f8e --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/acpbench/gen_2shot_with_pddl/prog.yaml @@ -0,0 +1,24 @@ +task: acp_prog_gen_with_pddl +dataset_name: acp_prog_gen +include: _gen_yaml_2shot +fewshot_config: + sampler: first_n + samples: + - context: "A robot is in a grid and can only move to places that are connected to its current position. \nThe grid size is 5x5, and the locations are of the form fi-jf (e.g., f3-2f or f0-1f). The grid cells are connected to their neighbors (e.g., f1-2f is connected to the four neighbors f0-2f, f2-2f, f1-1f, and f1-3f). Some positions on the grid are locked and can be opened with a key of a matching shape. The robot has an arm that can pick up a key when the key is in same location as the robot and the arm is empty. \nThere are 2 keys in 0 different shapes: Key key0-0 is of shape shape0, Key key0-1 is of shape shape0. \nCurrently, the robot is at position f2-2f and its arm is empty. All the positions are open except the following: f4-2f has shape0 shaped lock, f2-0f has shape0 shaped lock. Key key0-1 is at position f1-3f. Key key0-0 is at position f2-2f. The available propositions are: (at ?r ?x) - Key ?r is at ?x location, (at-robot ?x) - Robot is at ?x location, (locked ?x) - Location ?x is locked, (holding ?k) - Robot is holding ?k, (open ?x) - Location ?x is open, and (arm-empty) - Robot is not holding anything." + question: "Break down the outcomes of performing the action \"retrieve the key key0-0 from its current position f0-1f\" into two lists, positive effects and negative effects. Positive effects are the propositions that are false in the current state but will become true after performing the action. Negative effects are the propositions that are true in the current state and will become false after performing the action." + answer: "[(at-robot f1-2f)] [(at-robot f2-2f)]" + PDDL_domain: "(define (domain grid)\n (:requirements :strips :typing)\n (:types key place shape - object)\n (:predicates (arm-empty) (at ?r - key ?x - place) (at-robot ?x - place) (conn ?x - place ?y - place) (holding ?k - key) (key-shape ?k - key ?s - shape) (lock-shape ?x - place ?s - shape) (locked ?x - place) (open ?x - place))\n (:action move\n :parameters (?curpos - place ?nextpos - place)\n :precondition (and (at-robot ?curpos) (conn ?curpos ?nextpos) (open ?nextpos))\n :effect (and (at-robot ?nextpos) (not (at-robot ?curpos)))\n )\n (:action pickup\n :parameters (?curpos - place ?key - key)\n :precondition (and (at-robot ?curpos) (at ?key ?curpos) (arm-empty))\n :effect (and (holding ?key) (not (at ?key ?curpos)) (not (arm-empty)))\n )\n (:action pickup-and-loose\n :parameters (?curpos - place ?newkey - key ?oldkey - key)\n :precondition (and (at-robot ?curpos) (holding ?oldkey) (at ?newkey ?curpos))\n :effect (and (holding ?newkey) (at ?oldkey ?curpos) (not (holding ?oldkey)) (not (at ?newkey ?curpos)))\n )\n (:action putdown\n :parameters (?curpos - place ?key - key)\n :precondition (and (at-robot ?curpos) (holding ?key))\n :effect (and (arm-empty) (at ?key ?curpos) (not (holding ?key)))\n )\n (:action unlock\n :parameters (?curpos - place ?lockpos - place ?key - key ?shape - shape)\n :precondition (and (conn ?curpos ?lockpos) (key-shape ?key ?shape) (lock-shape ?lockpos ?shape) (at-robot ?curpos) (locked ?lockpos) (holding ?key))\n :effect (and (open ?lockpos) (not (locked ?lockpos)))\n )\n)" + PDDL_problem: "(define (problem grid-x5-y5-t1-k2-l2-p100)\n (:domain grid)\n (:requirements :strips :typing)\n (:objects key0-0 key0-1 - key f0-0f f0-1f f0-2f f0-3f f0-4f f1-0f f1-1f f1-2f f1-3f f1-4f f2-0f f2-1f f2-2f f2-3f f2-4f f3-0f f3-1f f3-2f f3-3f f3-4f f4-0f f4-1f f4-2f f4-3f f4-4f - place shape0 - shape)\n (:init (arm-empty) (at key0-0 f2-2f) (at key0-1 f1-3f) (at-robot f2-2f) (conn f0-0f f0-1f) (conn f0-0f f1-0f) (conn f0-1f f0-0f) (conn f0-1f f0-2f) (conn f0-1f f1-1f) (conn f0-2f f0-1f) (conn f0-2f f0-3f) (conn f0-2f f1-2f) (conn f0-3f f0-2f) (conn f0-3f f0-4f) (conn f0-3f f1-3f) (conn f0-4f f0-3f) (conn f0-4f f1-4f) (conn f1-0f f0-0f) (conn f1-0f f1-1f) (conn f1-0f f2-0f) (conn f1-1f f0-1f) (conn f1-1f f1-0f) (conn f1-1f f1-2f) (conn f1-1f f2-1f) (conn f1-2f f0-2f) (conn f1-2f f1-1f) (conn f1-2f f1-3f) (conn f1-2f f2-2f) (conn f1-3f f0-3f) (conn f1-3f f1-2f) (conn f1-3f f1-4f) (conn f1-3f f2-3f) (conn f1-4f f0-4f) (conn f1-4f f1-3f) (conn f1-4f f2-4f) (conn f2-0f f1-0f) (conn f2-0f f2-1f) (conn f2-0f f3-0f) (conn f2-1f f1-1f) (conn f2-1f f2-0f) (conn f2-1f f2-2f) (conn f2-1f f3-1f) (conn f2-2f f1-2f) (conn f2-2f f2-1f) (conn f2-2f f2-3f) (conn f2-2f f3-2f) (conn f2-3f f1-3f) (conn f2-3f f2-2f) (conn f2-3f f2-4f) (conn f2-3f f3-3f) (conn f2-4f f1-4f) (conn f2-4f f2-3f) (conn f2-4f f3-4f) (conn f3-0f f2-0f) (conn f3-0f f3-1f) (conn f3-0f f4-0f) (conn f3-1f f2-1f) (conn f3-1f f3-0f) (conn f3-1f f3-2f) (conn f3-1f f4-1f) (conn f3-2f f2-2f) (conn f3-2f f3-1f) (conn f3-2f f3-3f) (conn f3-2f f4-2f) (conn f3-3f f2-3f) (conn f3-3f f3-2f) (conn f3-3f f3-4f) (conn f3-3f f4-3f) (conn f3-4f f2-4f) (conn f3-4f f3-3f) (conn f3-4f f4-4f) (conn f4-0f f3-0f) (conn f4-0f f4-1f) (conn f4-1f f3-1f) (conn f4-1f f4-0f) (conn f4-1f f4-2f) (conn f4-2f f3-2f) (conn f4-2f f4-1f) (conn f4-2f f4-3f) (conn f4-3f f3-3f) (conn f4-3f f4-2f) (conn f4-3f f4-4f) (conn f4-4f f3-4f) (conn f4-4f f4-3f) (key-shape key0-0 shape0) (key-shape key0-1 shape0) (lock-shape f2-0f shape0) (lock-shape f4-2f shape0) (locked f2-0f) (locked f4-2f) (open f0-0f) (open f0-1f) (open f0-2f) (open f0-3f) (open f0-4f) (open f1-0f) (open f1-1f) (open f1-2f) (open f1-3f) (open f1-4f) (open f2-1f) (open f2-2f) (open f2-3f) (open f2-4f) (open f3-0f) (open f3-1f) (open f3-2f) (open f3-3f) (open f3-4f) (open f4-0f) (open f4-1f) (open f4-3f) (open f4-4f))\n (:goal (and (at key0-0 f2-0f) (at key0-1 f1-3f)))\n)" + - context: "There are several cities, each containing several locations, some of which are airports. There are also trucks, which can drive within a single city, and airplanes, which can fly between airports. The goal is to get some packages from various locations to various new locations. \nThere are 2 trucks and 1 airplane, as well as 4 packages. There are 4 locations across 2 cities. \nThe locations are in cities as follows: l0-0 and l0-1 are in c0; l1-0 and l1-1 are in c1. \nCurrently, p2 and t1 are at l1-0, a0 and t0 are at l0-0, p0 and p3 are in a0, p1 is in t1. The available propositions are: (at ?obj ?loc) - ?obj is at ?loc and (in ?obj1 ?obj2) - ?obj1 is in ?obj2." + question: "Break down the outcomes of performing the action \"navigate the truck t1 which is in location l1-0 in city c1 to another location l1-1 in the same city\" into two lists, positive effects and negative effects. Positive effects are the propositions that are false in the current state but will become true after performing the action. Negative effects are the propositions that are true in the current state and will become false after performing the action." + answer: "[(at t1 l1-1)] [(at t1 l1-0)]" + PDDL_domain: "(define (domain logistics-strips)\n (:requirements :strips :typing) \n\n (:types \n location locatable city - object \n package movable - locatable\n airport - location\n airplane truck - movable \n )\t\t\n \n (:predicates \t\n\t\t(at ?obj - locatable ?loc - location)\n\t\t(in ?obj1 - package ?obj2 - movable)\n\t\t(in-city ?obj - location ?city - city))\n\n\n(:action LOAD-TRUCK\n :parameters\n (?obj - package\n ?truck - truck\n ?loc - location)\n :precondition\n (and \n (at ?truck ?loc) (at ?obj ?loc))\n :effect\n (and (not (at ?obj ?loc)) (in ?obj ?truck)))\n\n(:action LOAD-AIRPLANE\n :parameters\n (?obj - package\n ?airplane - airplane\n ?loc - location)\n :precondition\n (and \n (at ?obj ?loc) (at ?airplane ?loc))\n :effect\n (and (not (at ?obj ?loc)) (in ?obj ?airplane)))\n\n\n\n(:action UNLOAD-TRUCK\n :parameters\n (?obj - package\n ?truck - truck\n ?loc - location)\n :precondition\n (and \n (at ?truck ?loc) (in ?obj ?truck))\n :effect\n (and (not (in ?obj ?truck)) (at ?obj ?loc)))\n\n(:action UNLOAD-AIRPLANE\n :parameters\n (?obj - package\n ?airplane - airplane\n ?loc - location)\n :precondition\n (and \n (in ?obj ?airplane) (at ?airplane ?loc))\n :effect\n (and (not (in ?obj ?airplane)) (at ?obj ?loc)))\n\n(:action DRIVE-TRUCK\n :parameters\n (?truck - truck\n ?loc-from - location\n ?loc-to - location\n ?city - city)\n :precondition\n (and \n (at ?truck ?loc-from)\n (in-city ?loc-from ?city)\n (in-city ?loc-to ?city))\n :effect\n (and (not (at ?truck ?loc-from)) (at ?truck ?loc-to)))\n\n(:action FLY-AIRPLANE\n :parameters\n (?airplane - airplane\n ?loc-from - airport\n ?loc-to - airport)\n :precondition\n (and \n\t(at ?airplane ?loc-from))\n :effect\n (and (not (at ?airplane ?loc-from)) (at ?airplane ?loc-to)))\n)" + PDDL_problem: "(define (problem logistics-c2-s2-p4-a1)\n (:domain logistics-strips)\n (:requirements :strips :typing)\n (:objects a0 - airplane l0-0 l1-0 - airport c0 c1 - city l0-1 l1-1 - location p0 p1 p2 p3 - package t0 t1 - truck)\n (:init (at a0 l0-0) (at p2 l1-0) (at t0 l0-0) (at t1 l1-0) (in p0 a0) (in p1 t1) (in p3 a0) (in-city l0-0 c0) (in-city l0-1 c0) (in-city l1-0 c1) (in-city l1-1 c1))\n (:goal (and (at p0 l0-0) (at p1 l1-0) (at p2 l1-0) (at p3 l0-1)))\n)" +doc_to_text: "# PDDL DOMAIN \n\n```\n{{PDDL_domain}}\n```\n\n# PDDL PROBLEM \n\n```\n{{PDDL_problem}}\n```\n\n**Question**: {{context}} {{question}} Provide only the two lists with the ground propositions. **Final Answer**:" +filter_list: + - name: "acp_grammar_parse" + filter: + - function: "ACP_grammar_filter" + grammar_task: "progression_list" + clean: "pos_neg" + - function: "take_first" diff --git a/lm-evaluation-harness/lm_eval/tasks/acpbench/gen_2shot_with_pddl/reach.yaml b/lm-evaluation-harness/lm_eval/tasks/acpbench/gen_2shot_with_pddl/reach.yaml new file mode 100644 index 0000000000000000000000000000000000000000..6cb78bbd836d6448911bce0d4b6357a9f43139a0 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/acpbench/gen_2shot_with_pddl/reach.yaml @@ -0,0 +1,23 @@ +task: acp_reach_gen_with_pddl +dataset_name: acp_reach_gen +include: _gen_yaml_2shot +fewshot_config: + sampler: first_n + samples: + - context: "A robot is in a grid and can only move to places that are connected to its current position. \nThe grid size is 5x5, and the locations are of the form fi-jf (e.g., f3-2f or f0-1f). The grid cells are connected to their neighbors (e.g., f1-2f is connected to the four neighbors f0-2f, f2-2f, f1-1f, and f1-3f). Some positions on the grid are locked and can be opened with a key of a matching shape. The robot has an arm that can pick up a key when the key is in same location as the robot and the arm is empty. \nThere are 2 keys in 0 different shapes: Key key0-1 is of shape shape0, Key key0-0 is of shape shape0. \nCurrently, the robot is at position f3-1f and its arm is empty. All the positions are open except the following: f2-0f has shape0 shaped lock. Key key0-1 is at position f1-3f. Key key0-0 is at position f3-1f. The available propositions are: (at ?r ?x) - Key ?r is at ?x location, (at-robot ?x) - Robot is at ?x location, (locked ?x) - Location ?x is locked, (holding ?k) - Robot is holding ?k, (open ?x) - Location ?x is open, and (arm-empty) - Robot's arm is empty." + question: "What proposition can never hold in any potentially reachable state?" + answer: "(locked f2-2f)" + PDDL_domain: "(define (domain grid)\n (:requirements :strips :typing)\n (:types key place shape - object)\n (:predicates (arm-empty) (at ?r - key ?x - place) (at-robot ?x - place) (conn ?x - place ?y - place) (holding ?k - key) (key-shape ?k - key ?s - shape) (lock-shape ?x - place ?s - shape) (locked ?x - place) (open ?x - place))\n (:action move\n :parameters (?curpos - place ?nextpos - place)\n :precondition (and (at-robot ?curpos) (conn ?curpos ?nextpos) (open ?nextpos))\n :effect (and (at-robot ?nextpos) (not (at-robot ?curpos)))\n )\n (:action pickup\n :parameters (?curpos - place ?key - key)\n :precondition (and (at-robot ?curpos) (at ?key ?curpos) (arm-empty))\n :effect (and (holding ?key) (not (at ?key ?curpos)) (not (arm-empty)))\n )\n (:action pickup-and-loose\n :parameters (?curpos - place ?newkey - key ?oldkey - key)\n :precondition (and (at-robot ?curpos) (holding ?oldkey) (at ?newkey ?curpos))\n :effect (and (holding ?newkey) (at ?oldkey ?curpos) (not (holding ?oldkey)) (not (at ?newkey ?curpos)))\n )\n (:action putdown\n :parameters (?curpos - place ?key - key)\n :precondition (and (at-robot ?curpos) (holding ?key))\n :effect (and (arm-empty) (at ?key ?curpos) (not (holding ?key)))\n )\n (:action unlock\n :parameters (?curpos - place ?lockpos - place ?key - key ?shape - shape)\n :precondition (and (conn ?curpos ?lockpos) (key-shape ?key ?shape) (lock-shape ?lockpos ?shape) (at-robot ?curpos) (locked ?lockpos) (holding ?key))\n :effect (and (open ?lockpos) (not (locked ?lockpos)))\n )\n)" + PDDL_problem: "(define (problem grid-x5-y5-t1-k2-l2-p100)\n (:domain grid)\n (:requirements :strips :typing)\n (:objects key0-0 key0-1 - key f0-0f f0-1f f0-2f f0-3f f0-4f f1-0f f1-1f f1-2f f1-3f f1-4f f2-0f f2-1f f2-2f f2-3f f2-4f f3-0f f3-1f f3-2f f3-3f f3-4f f4-0f f4-1f f4-2f f4-3f f4-4f - place shape0 - shape)\n (:init (arm-empty) (at key0-0 f3-1f) (at key0-1 f1-3f) (at-robot f3-1f) (conn f0-0f f0-1f) (conn f0-0f f1-0f) (conn f0-1f f0-0f) (conn f0-1f f0-2f) (conn f0-1f f1-1f) (conn f0-2f f0-1f) (conn f0-2f f0-3f) (conn f0-2f f1-2f) (conn f0-3f f0-2f) (conn f0-3f f0-4f) (conn f0-3f f1-3f) (conn f0-4f f0-3f) (conn f0-4f f1-4f) (conn f1-0f f0-0f) (conn f1-0f f1-1f) (conn f1-0f f2-0f) (conn f1-1f f0-1f) (conn f1-1f f1-0f) (conn f1-1f f1-2f) (conn f1-1f f2-1f) (conn f1-2f f0-2f) (conn f1-2f f1-1f) (conn f1-2f f1-3f) (conn f1-2f f2-2f) (conn f1-3f f0-3f) (conn f1-3f f1-2f) (conn f1-3f f1-4f) (conn f1-3f f2-3f) (conn f1-4f f0-4f) (conn f1-4f f1-3f) (conn f1-4f f2-4f) (conn f2-0f f1-0f) (conn f2-0f f2-1f) (conn f2-0f f3-0f) (conn f2-1f f1-1f) (conn f2-1f f2-0f) (conn f2-1f f2-2f) (conn f2-1f f3-1f) (conn f2-2f f1-2f) (conn f2-2f f2-1f) (conn f2-2f f2-3f) (conn f2-2f f3-2f) (conn f2-3f f1-3f) (conn f2-3f f2-2f) (conn f2-3f f2-4f) (conn f2-3f f3-3f) (conn f2-4f f1-4f) (conn f2-4f f2-3f) (conn f2-4f f3-4f) (conn f3-0f f2-0f) (conn f3-0f f3-1f) (conn f3-0f f4-0f) (conn f3-1f f2-1f) (conn f3-1f f3-0f) (conn f3-1f f3-2f) (conn f3-1f f4-1f) (conn f3-2f f2-2f) (conn f3-2f f3-1f) (conn f3-2f f3-3f) (conn f3-2f f4-2f) (conn f3-3f f2-3f) (conn f3-3f f3-2f) (conn f3-3f f3-4f) (conn f3-3f f4-3f) (conn f3-4f f2-4f) (conn f3-4f f3-3f) (conn f3-4f f4-4f) (conn f4-0f f3-0f) (conn f4-0f f4-1f) (conn f4-1f f3-1f) (conn f4-1f f4-0f) (conn f4-1f f4-2f) (conn f4-2f f3-2f) (conn f4-2f f4-1f) (conn f4-2f f4-3f) (conn f4-3f f3-3f) (conn f4-3f f4-2f) (conn f4-3f f4-4f) (conn f4-4f f3-4f) (conn f4-4f f4-3f) (key-shape key0-0 shape0) (key-shape key0-1 shape0) (lock-shape f2-0f shape0) (lock-shape f4-2f shape0) (locked f2-0f) (open f0-0f) (open f0-1f) (open f0-2f) (open f0-3f) (open f0-4f) (open f1-0f) (open f1-1f) (open f1-2f) (open f1-3f) (open f1-4f) (open f2-1f) (open f2-2f) (open f2-3f) (open f2-4f) (open f3-0f) (open f3-1f) (open f3-2f) (open f3-3f) (open f3-4f) (open f4-0f) (open f4-1f) (open f4-2f) (open f4-3f) (open f4-4f))\n (:goal (and (at key0-0 f2-0f) (at key0-1 f1-3f)))\n)" + - context: "There are several cities, each containing several locations, some of which are airports. There are also trucks, which can drive within a single city, and airplanes, which can fly between airports. The goal is to get some packages from various locations to various new locations. \nThere are 2 trucks and 1 airplane, as well as 4 packages. There are 4 locations across 2 cities. \nThe locations are in cities as follows: l1-0 and l1-1 are in c1; l0-1 and l0-0 are in c0. \nCurrently, t1, a0, and p2 are at l1-0, t0 and p0 are at l0-0, p1 is in t1, p3 is in t0. The available propositions are: (at ?obj ?loc) - ?obj is at ?loc and (in ?obj1 ?obj2) - ?obj1 is in ?obj2." + question: "What proposition can never hold in any potentially reachable state?" + answer: "(at t0 l1-0)" + PDDL_domain: "(define (domain logistics-strips)\n (:requirements :strips :typing) \n\n (:types \n location locatable city - object \n package movable - locatable\n airport - location\n airplane truck - movable \n )\t\t\n \n (:predicates \t\n\t\t(at ?obj - locatable ?loc - location)\n\t\t(in ?obj1 - package ?obj2 - movable)\n\t\t(in-city ?obj - location ?city - city))\n\n\n(:action LOAD-TRUCK\n :parameters\n (?obj - package\n ?truck - truck\n ?loc - location)\n :precondition\n (and \n (at ?truck ?loc) (at ?obj ?loc))\n :effect\n (and (not (at ?obj ?loc)) (in ?obj ?truck)))\n\n(:action LOAD-AIRPLANE\n :parameters\n (?obj - package\n ?airplane - airplane\n ?loc - location)\n :precondition\n (and \n (at ?obj ?loc) (at ?airplane ?loc))\n :effect\n (and (not (at ?obj ?loc)) (in ?obj ?airplane)))\n\n\n\n(:action UNLOAD-TRUCK\n :parameters\n (?obj - package\n ?truck - truck\n ?loc - location)\n :precondition\n (and \n (at ?truck ?loc) (in ?obj ?truck))\n :effect\n (and (not (in ?obj ?truck)) (at ?obj ?loc)))\n\n(:action UNLOAD-AIRPLANE\n :parameters\n (?obj - package\n ?airplane - airplane\n ?loc - location)\n :precondition\n (and \n (in ?obj ?airplane) (at ?airplane ?loc))\n :effect\n (and (not (in ?obj ?airplane)) (at ?obj ?loc)))\n\n(:action DRIVE-TRUCK\n :parameters\n (?truck - truck\n ?loc-from - location\n ?loc-to - location\n ?city - city)\n :precondition\n (and \n (at ?truck ?loc-from)\n (in-city ?loc-from ?city)\n (in-city ?loc-to ?city))\n :effect\n (and (not (at ?truck ?loc-from)) (at ?truck ?loc-to)))\n\n(:action FLY-AIRPLANE\n :parameters\n (?airplane - airplane\n ?loc-from - airport\n ?loc-to - airport)\n :precondition\n (and \n\t(at ?airplane ?loc-from))\n :effect\n (and (not (at ?airplane ?loc-from)) (at ?airplane ?loc-to)))\n)" + PDDL_problem: "(define (problem logistics-c2-s2-p4-a1)\n (:domain logistics-strips)\n (:requirements :strips :typing)\n (:objects a0 - airplane l0-0 l1-0 - airport c0 c1 - city l0-1 l1-1 - location p0 p1 p2 p3 - package t0 t1 - truck)\n (:init (at a0 l1-0) (at p0 l0-0) (at p2 l1-0) (at t0 l0-0) (at t1 l1-0) (in p1 t1) (in p3 t0) (in-city l0-0 c0) (in-city l0-1 c0) (in-city l1-0 c1) (in-city l1-1 c1))\n (:goal (and (at p0 l0-0) (at p1 l1-0) (at p2 l1-0) (at p3 l0-1)))\n)" +doc_to_text: "# PDDL DOMAIN \n\n```\n{{PDDL_domain}}\n```\n\n# PDDL PROBLEM \n\n```\n{{PDDL_problem}}\n```\n\n**Question**: {{context}} {{question}} Provide one proposition or None. **Final Answer**:" +filter_list: + - name: "acp_grammar_parse" + filter: + - function: "ACP_grammar_filter" + grammar_task: "act" + - function: "take_first" diff --git a/lm-evaluation-harness/lm_eval/tasks/acpbench/gen_2shot_with_pddl/val.yaml b/lm-evaluation-harness/lm_eval/tasks/acpbench/gen_2shot_with_pddl/val.yaml new file mode 100644 index 0000000000000000000000000000000000000000..6012ebf92a546a200ab42987e8406dacd11c75e4 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/acpbench/gen_2shot_with_pddl/val.yaml @@ -0,0 +1,23 @@ +task: acp_val_gen_with_pddl +dataset_name: acp_val_gen +include: _gen_yaml_2shot +fewshot_config: + sampler: first_n + samples: + - context: "A robot is in a grid and can only move to places that are connected to its current position. \nThe grid size is 5x5, and the locations are of the form fi-jf (e.g., f3-2f or f0-1f). The grid cells are connected to their neighbors (e.g., f1-2f is connected to the four neighbors f0-2f, f2-2f, f1-1f, and f1-3f). Some positions on the grid are locked and can be opened with a key of a matching shape. The robot has an arm that can pick up a key when the key is in same location as the robot and the arm is empty. \nThere are 2 keys in 0 different shapes: Key key0-1 is of shape shape0, Key key0-0 is of shape shape0. \nCurrently, the robot is at position f3-3f and its arm is empty. All the positions are open except the following: f2-0f has shape0 shaped lock, f4-2f has shape0 shaped lock. Key key0-1 is at position f1-3f. Key key0-0 is at position f2-2f. The goal is to reach a state where the following facts hold: Key key0-1 is at f1-3f location and Key key0-0 is at f2-0f location. The available actions are: (unlock ?curpos ?lockpos ?key ?shape) - unlock the place ?lockpos with the key ?key of the shape ?shape from the current position place ?curpos, (move ?curpos ?nextpos) - move from place ?curpos to place ?nextpos, (pickup ?curpos ?key) - retrieve the key ?key from its current position ?curpos, (pickup-and-loose ?curpos ?newkey ?oldkey) - pick up the key ?newkey at the current position place ?curpos and loose the key ?oldkey being held, and (putdown ?curpos ?key) - put the key ?key at the current position place ?curpos." + question: "What is the first inapplicable action in the next sequence of actions: \"(unlock f1-0f f2-0f key0-1 shape0) (move f2-3f f2-2f) (pickup f2-2f key0-0) (move f2-2f f2-1f) (unlock f2-1f f2-0f key0-0 shape0) (move f2-1f f2-0f) (putdown f2-0f key0-0)\"?" + answer: "0" + PDDL_domain: "(define (domain grid)\n (:requirements :strips :typing)\n (:types key place shape - object)\n (:predicates (arm-empty) (at ?r - key ?x - place) (at-robot ?x - place) (conn ?x - place ?y - place) (holding ?k - key) (key-shape ?k - key ?s - shape) (lock-shape ?x - place ?s - shape) (locked ?x - place) (open ?x - place))\n (:action move\n :parameters (?curpos - place ?nextpos - place)\n :precondition (and (at-robot ?curpos) (conn ?curpos ?nextpos) (open ?nextpos))\n :effect (and (at-robot ?nextpos) (not (at-robot ?curpos)))\n )\n (:action pickup\n :parameters (?curpos - place ?key - key)\n :precondition (and (at-robot ?curpos) (at ?key ?curpos) (arm-empty))\n :effect (and (holding ?key) (not (at ?key ?curpos)) (not (arm-empty)))\n )\n (:action pickup-and-loose\n :parameters (?curpos - place ?newkey - key ?oldkey - key)\n :precondition (and (at-robot ?curpos) (holding ?oldkey) (at ?newkey ?curpos))\n :effect (and (holding ?newkey) (at ?oldkey ?curpos) (not (holding ?oldkey)) (not (at ?newkey ?curpos)))\n )\n (:action putdown\n :parameters (?curpos - place ?key - key)\n :precondition (and (at-robot ?curpos) (holding ?key))\n :effect (and (arm-empty) (at ?key ?curpos) (not (holding ?key)))\n )\n (:action unlock\n :parameters (?curpos - place ?lockpos - place ?key - key ?shape - shape)\n :precondition (and (conn ?curpos ?lockpos) (key-shape ?key ?shape) (lock-shape ?lockpos ?shape) (at-robot ?curpos) (locked ?lockpos) (holding ?key))\n :effect (and (open ?lockpos) (not (locked ?lockpos)))\n )\n)" + PDDL_problem: "(define (problem grid-x5-y5-t1-k2-l2-p100)\n (:domain grid)\n (:requirements :strips :typing)\n (:objects key0-0 key0-1 - key f0-0f f0-1f f0-2f f0-3f f0-4f f1-0f f1-1f f1-2f f1-3f f1-4f f2-0f f2-1f f2-2f f2-3f f2-4f f3-0f f3-1f f3-2f f3-3f f3-4f f4-0f f4-1f f4-2f f4-3f f4-4f - place shape0 - shape)\n (:init (arm-empty) (at key0-0 f2-2f) (at key0-1 f1-3f) (at-robot f3-3f) (conn f0-0f f0-1f) (conn f0-0f f1-0f) (conn f0-1f f0-0f) (conn f0-1f f0-2f) (conn f0-1f f1-1f) (conn f0-2f f0-1f) (conn f0-2f f0-3f) (conn f0-2f f1-2f) (conn f0-3f f0-2f) (conn f0-3f f0-4f) (conn f0-3f f1-3f) (conn f0-4f f0-3f) (conn f0-4f f1-4f) (conn f1-0f f0-0f) (conn f1-0f f1-1f) (conn f1-0f f2-0f) (conn f1-1f f0-1f) (conn f1-1f f1-0f) (conn f1-1f f1-2f) (conn f1-1f f2-1f) (conn f1-2f f0-2f) (conn f1-2f f1-1f) (conn f1-2f f1-3f) (conn f1-2f f2-2f) (conn f1-3f f0-3f) (conn f1-3f f1-2f) (conn f1-3f f1-4f) (conn f1-3f f2-3f) (conn f1-4f f0-4f) (conn f1-4f f1-3f) (conn f1-4f f2-4f) (conn f2-0f f1-0f) (conn f2-0f f2-1f) (conn f2-0f f3-0f) (conn f2-1f f1-1f) (conn f2-1f f2-0f) (conn f2-1f f2-2f) (conn f2-1f f3-1f) (conn f2-2f f1-2f) (conn f2-2f f2-1f) (conn f2-2f f2-3f) (conn f2-2f f3-2f) (conn f2-3f f1-3f) (conn f2-3f f2-2f) (conn f2-3f f2-4f) (conn f2-3f f3-3f) (conn f2-4f f1-4f) (conn f2-4f f2-3f) (conn f2-4f f3-4f) (conn f3-0f f2-0f) (conn f3-0f f3-1f) (conn f3-0f f4-0f) (conn f3-1f f2-1f) (conn f3-1f f3-0f) (conn f3-1f f3-2f) (conn f3-1f f4-1f) (conn f3-2f f2-2f) (conn f3-2f f3-1f) (conn f3-2f f3-3f) (conn f3-2f f4-2f) (conn f3-3f f2-3f) (conn f3-3f f3-2f) (conn f3-3f f3-4f) (conn f3-3f f4-3f) (conn f3-4f f2-4f) (conn f3-4f f3-3f) (conn f3-4f f4-4f) (conn f4-0f f3-0f) (conn f4-0f f4-1f) (conn f4-1f f3-1f) (conn f4-1f f4-0f) (conn f4-1f f4-2f) (conn f4-2f f3-2f) (conn f4-2f f4-1f) (conn f4-2f f4-3f) (conn f4-3f f3-3f) (conn f4-3f f4-2f) (conn f4-3f f4-4f) (conn f4-4f f3-4f) (conn f4-4f f4-3f) (key-shape key0-0 shape0) (key-shape key0-1 shape0) (lock-shape f2-0f shape0) (lock-shape f4-2f shape0) (locked f2-0f) (locked f4-2f) (open f0-0f) (open f0-1f) (open f0-2f) (open f0-3f) (open f0-4f) (open f1-0f) (open f1-1f) (open f1-2f) (open f1-3f) (open f1-4f) (open f2-1f) (open f2-2f) (open f2-3f) (open f2-4f) (open f3-0f) (open f3-1f) (open f3-2f) (open f3-3f) (open f3-4f) (open f4-0f) (open f4-1f) (open f4-3f) (open f4-4f))\n (:goal (and (at key0-0 f2-0f) (at key0-1 f1-3f)))\n)" + - context: "There are several cities, each containing several locations, some of which are airports. There are also trucks, which can drive within a single city, and airplanes, which can fly between airports. The goal is to get some packages from various locations to various new locations. \nThere are 2 trucks and 1 airplane, as well as 4 packages. There are 4 locations across 2 cities. \nThe locations are in cities as follows: l1-1 and l1-0 are in c1; l0-1 and l0-0 are in c0. \nCurrently, t1 and p0 are at l1-1, p3, p2, and p1 are at l1-0, t0 is at l0-1, a0 is at l0-0. The goal is to reach a state where the following facts hold: p0 is at l0-0, p3 is at l0-1, p2 is at l1-0, and p1 is at l1-0. The available actions are: (load-truck ?obj ?truck ?loc) - load object ?obj into truck ?truck at location ?loc, (load-airplane ?obj ?airplane ?loc) - load the object ?obj from location ?loc into the airplane ?airplane, (unload-truck ?obj ?truck ?loc) - offload the object ?obj from the truck ?truck at location ?loc, (unload-airplane ?obj ?airplane ?loc) - remove the object ?obj from the airplane ?airplane and place it on the location ?loc, (drive-truck ?truck ?loc-from ?loc-to ?city) - navigate the truck ?truck which is in location ?loc-from in city ?city to another location ?loc-to in the same city, and (fly-airplane ?airplane ?loc-from ?loc-to) - fly airplane ?airplane from airport ?loc-from to airport ?loc-to." + question: "What is the first inapplicable action in the next sequence of actions: \"(drive-truck t0 l0-1 l0-0 c0) (fly-airplane a0 l0-0 l1-0) (load-airplane p3 a0 l1-0) (load-truck p0 t1 l1-1) (drive-truck t1 l1-1 l1-0 c1) (unload-truck p0 t1 l1-0) (load-airplane p0 a0 l1-0) (fly-airplane a0 l1-0 l0-0) (unload-airplane p0 a0 l0-0) (unload-airplane p3 a0 l0-0) (load-truck p3 t0 l0-0) (drive-truck t0 l0-0 l0-1 c0) (unload-airplane p3 a0 l0-0)\"?" + answer: "12" + PDDL_domain: "(define (domain logistics-strips)\n (:requirements :strips :typing) \n\n (:types \n location locatable city - object \n package movable - locatable\n airport - location\n airplane truck - movable \n )\t\t\n \n (:predicates \t\n\t\t(at ?obj - locatable ?loc - location)\n\t\t(in ?obj1 - package ?obj2 - movable)\n\t\t(in-city ?obj - location ?city - city))\n\n\n(:action LOAD-TRUCK\n :parameters\n (?obj - package\n ?truck - truck\n ?loc - location)\n :precondition\n (and \n (at ?truck ?loc) (at ?obj ?loc))\n :effect\n (and (not (at ?obj ?loc)) (in ?obj ?truck)))\n\n(:action LOAD-AIRPLANE\n :parameters\n (?obj - package\n ?airplane - airplane\n ?loc - location)\n :precondition\n (and \n (at ?obj ?loc) (at ?airplane ?loc))\n :effect\n (and (not (at ?obj ?loc)) (in ?obj ?airplane)))\n\n\n\n(:action UNLOAD-TRUCK\n :parameters\n (?obj - package\n ?truck - truck\n ?loc - location)\n :precondition\n (and \n (at ?truck ?loc) (in ?obj ?truck))\n :effect\n (and (not (in ?obj ?truck)) (at ?obj ?loc)))\n\n(:action UNLOAD-AIRPLANE\n :parameters\n (?obj - package\n ?airplane - airplane\n ?loc - location)\n :precondition\n (and \n (in ?obj ?airplane) (at ?airplane ?loc))\n :effect\n (and (not (in ?obj ?airplane)) (at ?obj ?loc)))\n\n(:action DRIVE-TRUCK\n :parameters\n (?truck - truck\n ?loc-from - location\n ?loc-to - location\n ?city - city)\n :precondition\n (and \n (at ?truck ?loc-from)\n (in-city ?loc-from ?city)\n (in-city ?loc-to ?city))\n :effect\n (and (not (at ?truck ?loc-from)) (at ?truck ?loc-to)))\n\n(:action FLY-AIRPLANE\n :parameters\n (?airplane - airplane\n ?loc-from - airport\n ?loc-to - airport)\n :precondition\n (and \n\t(at ?airplane ?loc-from))\n :effect\n (and (not (at ?airplane ?loc-from)) (at ?airplane ?loc-to)))\n)" + PDDL_problem: "(define (problem logistics-c2-s2-p4-a1)\n (:domain logistics-strips)\n (:requirements :strips :typing)\n (:objects a0 - airplane l0-0 l1-0 - airport c0 c1 - city l0-1 l1-1 - location p0 p1 p2 p3 - package t0 t1 - truck)\n (:init (at a0 l0-0) (at p0 l1-1) (at p1 l1-0) (at p2 l1-0) (at p3 l1-0) (at t0 l0-1) (at t1 l1-1) (in-city l0-0 c0) (in-city l0-1 c0) (in-city l1-0 c1) (in-city l1-1 c1))\n (:goal (and (at p0 l0-0) (at p1 l1-0) (at p2 l1-0) (at p3 l0-1)))\n)" +doc_to_text: "# PDDL DOMAIN \n\n```\n{{PDDL_domain}}\n```\n\n# PDDL PROBLEM \n\n```\n{{PDDL_problem}}\n```\n\n**Question**: {{context}} {{question}} Provide only the index of the action. **Final Answer**:" +filter_list: + - name: "acp_grammar_parse" + filter: + - function: "ACP_grammar_filter" + grammar_task: "index" + - function: "take_first" diff --git a/lm-evaluation-harness/lm_eval/tasks/acpbench/mcq_cot_2shot/_mcq_cot_2shot_yaml b/lm-evaluation-harness/lm_eval/tasks/acpbench/mcq_cot_2shot/_mcq_cot_2shot_yaml new file mode 100644 index 0000000000000000000000000000000000000000..d8ace1f60098ff1a576d1ca91b4d3465a33ee9e8 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/acpbench/mcq_cot_2shot/_mcq_cot_2shot_yaml @@ -0,0 +1,39 @@ +tag: + - acp_mcq_cot_2shot + - acp_bench +output_type: generate_until +dataset_path: ibm-research/acp_bench +test_split: test +num_fewshot: 2 +doc_to_target: "{{answer}}" +doc_to_text: "**Question**: {{context}} {{question}} **Thoughts**:" +generation_kwargs: + until: + - "\n\n\n\n" + - "**Question**:" + - "**Question:**" + - "Q:" + do_sample: false + temperature: 0.0 + max_gen_toks: 1024 +metric_list: + - metric: exact_match + ignore_case: true + ignore_punctuation: true + regexes_to_ignore: + - "\\.$" + - "," + - "\\\\" + - "\n" + - '"' +filter_list: + - name: "mcq-extract" + filter: + - function: multi_choice_regex + group_select: -1 + ignore_case: true + ignore_punctuation: true + regex_pattern: '(((?<=[answer is ])[A-D])|([A-D]\n)|([A-D]\.)|( [A-D] )|(^[A-D]$)|(\[[A-D]\])|([A-D])|(?<=..Final Answer..: )(.*)(?=.)|(?<=..answer..: )(.*)(?=.)|(?<=..Answer..: )(.*)(?=.))' + - function: "take_first" +metadata: + version: 1.0 diff --git a/lm-evaluation-harness/lm_eval/tasks/acpbench/mcq_cot_2shot/act_reach.yaml b/lm-evaluation-harness/lm_eval/tasks/acpbench/mcq_cot_2shot/act_reach.yaml new file mode 100644 index 0000000000000000000000000000000000000000..bde192e4d68bc87e5e877fd6f35bb0bb9ed52ebd --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/acpbench/mcq_cot_2shot/act_reach.yaml @@ -0,0 +1,12 @@ +task: acp_areach_mcq +dataset_name: acp_areach_mcq +include: _mcq_cot_2shot_yaml +fewshot_config: + sampler: first_n + samples: + - context: 'This is a ferry domain, where the task is to transport cars from their start to their goal locations, using a ferry. Each location is accessible by ferry from each other location. The cars can be debarked or boarded, and the ferry can carry only one car at a time. There are 2 locations and 2 cars, numbered consecutively. Currently, the ferry is at l0, with the car c1 on board. The cars are at locations as follows: c0 is at l0.' + question: 'Which of the following actions can eventually be applied? A. embark the car c0 at location l0 on to the ferry. B. travel by sea from location c0 to location c1. C. fly from location l0 to location l1. D. board the car c0 at location l0 into the airplane.' + answer: "Let's think step by step. Step 1: Verify if there is a sequence of actions which transforms the current state into a state where the precondition of the action \"embark the car c0 at location l0 on to the ferry\" hold. Step 2: The following sequence of actions would transition to such a state: travel by sea from location l0 to location l1, debark the car c1 from the ferry to location l1, travel by sea from location l1 to location l0. **Final Answer**: A." + - context: 'There are several cities, each containing several locations, some of which are airports. There are also trucks, which can drive within a single city, and airplanes, which can fly between airports. The goal is to get some packages from various locations to various new locations. There are 2 trucks and 1 airplane, as well as 4 packages. There are 4 locations across 2 cities. The locations are in cities as follows: l1-1 and l1-0 are in c1; l0-0 and l0-1 are in c0. Currently, a0 is at l0-0, t0 is at l0-1, p3 and t1 are at l1-0, p0 is in t1, p2 and p1 are in a0.' + question: 'Which of the following actions can eventually be applied? A. offload the object p3 from the truck p3 at location l1-1. B. navigate the truck c1 which is in location p0 in city l0-1 to another location t1 in the same city. C. fly the airplane a0 from airport l1-0 to airport l0-0. D. fly the airplane a0 to the airport l0-0 in city l1-0.' + answer: "Let's think step by step. Step 1: Verify if there is a sequence of actions which transforms the current state into a state where the precondition of the action \"fly the airplane a0 from airport l1-0 to airport l0-0\" hold. Step 2: The following sequence of actions would transition to such a state: drive truck t0 from location l0-1 in city c0 to location l0-0 in the same city, fly the airplane a0 from location l0-0 to location l1-0, navigate the truck t1 which is in location l1-0 in city c1 to another location l1-1 in the same city. **Final Answer**: C." diff --git a/lm-evaluation-harness/lm_eval/tasks/acpbench/mcq_cot_2shot/app.yaml b/lm-evaluation-harness/lm_eval/tasks/acpbench/mcq_cot_2shot/app.yaml new file mode 100644 index 0000000000000000000000000000000000000000..38defabf9d5a254d25af602e56a3e54d1077e395 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/acpbench/mcq_cot_2shot/app.yaml @@ -0,0 +1,12 @@ +task: acp_app_mcq +dataset_name: acp_app_mcq +include: _mcq_cot_2shot_yaml +fewshot_config: + sampler: first_n + samples: + - context: 'This is a ferry domain, where the task is to transport cars from their start to their goal locations, using a ferry. Each location is accessible by ferry from each other location. The cars can be debarked or boarded, and the ferry can carry only one car at a time. There are 2 locations and 2 cars, numbered consecutively. Currently, the ferry is at l1 location and it is empty. The cars are at locations as follows: c1 and c0 are at l0.' + question: 'Which of the following actions will be applicable in this state? A. unload the car c1 from the ferry to location l1. B. load the car c0 at location l1 on to the ferry. C. load the car c0 at location l0 on to the ferry. D. sail from location l1 to location l0.' + answer: "Let's think step by step. Step 1: In order to apply the action \"sail from location l1 to location l0\", the following fact(s) must hold in this state: The ferry is at l1 location Step 2: These facts hold in the mentioned state, so the action \"sail from location l1 to location l0\" is applicable. **Final Answer**: D." + - context: 'There are several cities, each containing several locations, some of which are airports. There are also trucks, which can drive within a single city, and airplanes, which can fly between airports. The goal is to get some packages from various locations to various new locations. There are 2 trucks and 1 airplane, as well as 4 packages. There are 4 locations across 2 cities. The locations are in cities as follows: l1-1 and l1-0 are in c1; l0-0 and l0-1 are in c0. Currently, a0 and t0 are at l0-0, t1 is at l1-1, p2 is at l1-0, p1 and p3 are in t1, p0 is in t0.' + question: 'Which of the following actions will be applicable in this state? A. load object p1 into airplane a0 at location l1-0. B. unload the object p2 from the airplane a0 at location l1-0. C. navigate the truck t1 from location l1-1 in city c1 to location l1-1 in the same city. D. operate the airplane a0 from airport l1-0 to airport l0-0.' + answer: "Let's think step by step. Step 1: In order to apply the action \"navigate the truck t1 from location l1-1 in city c1 to location l1-1 in the same city\", the following fact(s) must hold in this state: t1 is at l1-1 Step 2: These facts hold in the mentioned state, so the action \"navigate the truck t1 from location l1-1 in city c1 to location l1-1 in the same city\" is applicable. **Final Answer**: C." diff --git a/lm-evaluation-harness/lm_eval/tasks/acpbench/mcq_cot_2shot/land.yaml b/lm-evaluation-harness/lm_eval/tasks/acpbench/mcq_cot_2shot/land.yaml new file mode 100644 index 0000000000000000000000000000000000000000..af0df12032949c8d7cde85765f7f68e50ff0f555 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/acpbench/mcq_cot_2shot/land.yaml @@ -0,0 +1,12 @@ +task: acp_land_mcq +dataset_name: acp_land_mcq +include: _mcq_cot_2shot_yaml +fewshot_config: + sampler: first_n + samples: + - context: 'This is a ferry domain, where the task is to transport cars from their start to their goal locations, using a ferry. Each location is accessible by ferry from each other location. The cars can be debarked or boarded, and the ferry can carry only one car at a time. There are 2 locations and 2 cars, numbered consecutively. Currently, the ferry is at l0 location and it is empty. The cars are at locations as follows: c1 is at l0; c0 is at l1. The goal is to reach a state where the following facts hold: Car c1 is at location l1 and Car c0 is at location l1.' + question: 'Which of the following facts is a landmark (must hold at some point along any plan) for the current state? A. Car c0 is on the ferry. B. Ferry has car c1 on board and Car c0 is at location l0. C. Ferry has car c1 on board. D. Ferry has car c1 on board and Car c0 is on the ferry.' + answer: "Let's think step by step. Step 1: A fact is a landmark if it must hold at some point along any plan. Step 2: The fact \"Ferry has car c1 on board\" can be found by a simple procedure that traces back such atoms from the goal. **Final Answer**: C." + - context: 'There are several cities, each containing several locations, some of which are airports. There are also trucks, which can drive within a single city, and airplanes, which can fly between airports. The goal is to get some packages from various locations to various new locations. There are 2 trucks and 1 airplane, as well as 4 packages. There are 4 locations across 2 cities. The locations are in cities as follows: l1-1 and l1-0 are in c1; l0-1 and l0-0 are in c0. Currently, p1 and t0 are at l0-0, a0, p0, p3, and p2 are at l1-0, t1 is at l1-1. The goal is to reach a state where the following facts hold: p0 is at l0-0, p1 is at l1-0, p3 is at l0-1, and p2 is at l1-0.' + question: 'Which of the following facts is a landmark (must hold at some point along any plan) for the current state? A. p3 is at l0-0. B. p1 is at l0-1. C. p0 is in t0. D. p2 is in a0.' + answer: "Let's think step by step. Step 1: A fact is a landmark if it must hold at some point along any plan. Step 2: The fact \"p3 is at l0-0\" can be found by a simple procedure that traces back such atoms from the goal. **Final Answer**: A." diff --git a/lm-evaluation-harness/lm_eval/tasks/acpbench/mcq_cot_2shot/prog.yaml b/lm-evaluation-harness/lm_eval/tasks/acpbench/mcq_cot_2shot/prog.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c840654f4130e5bb8d47b470445b2318cd8cff44 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/acpbench/mcq_cot_2shot/prog.yaml @@ -0,0 +1,12 @@ +task: acp_prog_mcq +dataset_name: acp_prog_mcq +include: _mcq_cot_2shot_yaml +fewshot_config: + sampler: first_n + samples: + - context: 'This is a ferry domain, where the task is to transport cars from their start to their goal locations, using a ferry. Each location is accessible by ferry from each other location. The cars can be debarked or boarded, and the ferry can carry only one car at a time. There are 2 locations and 2 cars, numbered consecutively. Currently, the ferry is at l1, with the car c1 on board. The cars are at locations as follows: c0 is at l1.' + question: 'Which the following facts hold after performing the action \"travel by sea from location l1 to location l0\" in the current state? **Possible Answers**: A. Car c0 is at location l1 and The ferry is at l1 location. B. The ferry is at l0 location and The ferry is at l1 location. C. The ferry is at l0 location. D. The ferry is at l0 location and Car c0 is at location l1.' + answer: "Let's think step by step. Step 1: The following fact(s) do not hold in the current state: The ferry is at l0 location. Step 2: The action adds the following fact(s): The ferry is at l0 location Step 3: The following fact(s) hold in the current state: Car c0 is at location l1. Step 4: The action deletes the following fact(s): The ferry is at l1 location Step 5: Fact(s) \"The ferry is at l0 location\" are added and Fact(s) \"Car c0 is at location l1\" are not deleted. **Final Answer**: D." + - context: 'There are several cities, each containing several locations, some of which are airports. There are also trucks, which can drive within a single city, and airplanes, which can fly between airports. The goal is to get some packages from various locations to various new locations. There are 2 trucks and 1 airplane, as well as 4 packages. There are 4 locations across 2 cities. The locations are in cities as follows: l1-1 and l1-0 are in c1; l0-1 and l0-0 are in c0. Currently, a0 is at l0-0, t1 and p0 are at l1-1, t0 is at l0-1, p1 is in t1, p2 and p3 are in a0.' + question: 'Which the following facts hold after performing the action \"drive truck t0 from location l0-1 in city c0 to location l0-1 in the same city\" in the current state? A. p3 is in t1. B. a0 is at l0-0 and p3 is in t1. C. a0 is at l0-0. D. None of the above.' + answer: "Let's think step by step. Step 1: The following fact(s) hold in the current state: a0 is at l0-0. Step 2: The action deletes the following fact(s): t0 is at l0-1 Step 3: Fact(s) \"a0 is at l0-0\" are not deleted. **Final Answer**: C." diff --git a/lm-evaluation-harness/lm_eval/tasks/acpbench/mcq_cot_2shot/reach.yaml b/lm-evaluation-harness/lm_eval/tasks/acpbench/mcq_cot_2shot/reach.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c13ff3d862e20102642aea08b90eb5880afdf223 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/acpbench/mcq_cot_2shot/reach.yaml @@ -0,0 +1,12 @@ +task: acp_reach_mcq +dataset_name: acp_reach_mcq +include: _mcq_cot_2shot_yaml +fewshot_config: + sampler: first_n + samples: + - context: 'This is a ferry domain, where the task is to transport cars from their start to their goal locations, using a ferry. Each location is accessible by ferry from each other location. The cars can be debarked or boarded, and the ferry can carry only one car at a time. There are 2 locations and 2 cars, numbered consecutively. Currently, the ferry is at l0, with the car c0 on board. The cars are at locations as follows: c1 is at l1.' + question: 'Which of the following options can hold in a state that can potentially be reached? A. There are no cars on the ferry and The ferry is at l1 location. B. Car l1 is at location c1. C. Ferry has car c0 on board and There are no cars on the ferry. D. The ferry is at c0 location and Car c1 is at location l1.' + answer: "Let's think step by step. Step 1: Verify if the following facts hold in the current state: There are no cars on the ferry and The ferry is at l1 location. Step 2: These facts do not hold. Step 3: Verify if there is a sequence of actions which transforms the current state into a state where these facts hold. Step 4: The following sequence of actions would transition to such a state: debark the car c0 from the ferry to location l0, sail from location l0 to location l1. **Final Answer**: A." + - context: 'There are several cities, each containing several locations, some of which are airports. There are also trucks, which can drive within a single city, and airplanes, which can fly between airports. The goal is to get some packages from various locations to various new locations. There are 2 trucks and 1 airplane, as well as 4 packages. There are 4 locations across 2 cities. The locations are in cities as follows: l1-1 and l1-0 are in c1; l0-0 and l0-1 are in c0. Currently, p1 and t1 are at l1-1, a0 and p0 are at l0-0, t0 and p3 are at l0-1, p2 is at l1-0.' + question: 'Which of the following options can hold in a state that can potentially be reached? A. p2 is at p0. B. t0 is at l0-0. C. p3 is in t1 and p3 is in t0. D. l1-1 is at p1.' + answer: "Let's think step by step. Step 1: Verify if the following fact holds in the current state: t0 is at l0-0. Step 2: The fact does not hold. Step 3: Verify if there is a sequence of actions which transforms the current state into a state where the fact holds. Step 4: The following sequence of actions would transition to such a state: drive the truck t0 in city c0 from location l0-1 to location l0-0, navigate the truck t0 which is in location l0-0 in city c0 to another location l0-0 in the same city. **Final Answer**: B." diff --git a/lm-evaluation-harness/lm_eval/tasks/acpbench/mcq_cot_2shot/val.yaml b/lm-evaluation-harness/lm_eval/tasks/acpbench/mcq_cot_2shot/val.yaml new file mode 100644 index 0000000000000000000000000000000000000000..7aecbc7d5d8e80afab8c05800ef9d507e4cd632a --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/acpbench/mcq_cot_2shot/val.yaml @@ -0,0 +1,12 @@ +task: acp_val_mcq +dataset_name: acp_val_mcq +include: _mcq_cot_2shot_yaml +fewshot_config: + sampler: first_n + samples: + - context: 'This is a ferry domain, where the task is to transport cars from their start to their goal locations, using a ferry. Each location is accessible by ferry from each other location. The cars can be debarked or boarded, and the ferry can carry only one car at a time. There are 2 locations and 2 cars, numbered consecutively. Currently, the ferry is at l0 location and it is empty. The cars are at locations as follows: c1 and c0 are at l0. The goal is to reach a state where the following facts hold: Car c0 is at location l1 and Car c1 is at location l1.' + question: 'Which of the following claims is true with regard to the following sequence of actions \"board the car c0 at the location l0, travel by sea from location l0 to location l1, unload the car c0 from the ferry to location l1, travel by sea from location l1 to location l0, board the car c1 at location l0, sail from location l0 to location l1, debark the car c1 from the ferry to location l1\" and the current state? A. The sequence is not applicable. B. The sequence is a plan. C. The sequence is applicable, but does not achieve the goal. D. The sequence is not valid.' + answer: "Let's think step by step. Step 1: For a sequence of actions to be a plan, all actions should be valid, applicable in sequence, and achieve the goal. Step 2: The action sequence is applicable and it achieves the goal. **Final Answer**: B." + - context: 'There are several cities, each containing several locations, some of which are airports. There are also trucks, which can drive within a single city, and airplanes, which can fly between airports. The goal is to get some packages from various locations to various new locations. There are 3 trucks and 1 airplane, as well as 4 packages. There are 9 locations across 3 cities. The locations are in cities as follows: l1-2, l1-0, and l1-1 are in c1; l0-0, l0-1, and l0-2 are in c0; l2-1, l2-2, and l2-0 are in c2. Currently, p2 and t1 are at l1-2, p3 is at l2-0, t0 and p0 are at l0-2, p1 is at l1-0, a0 is at l0-0, t2 is at l2-2. The goal is to reach a state where the following facts hold: p1 is at l1-0, p3 is at l2-0, p2 is at l0-1, and p0 is at l1-2.' + question: 'Which of the following claims is true with regard to the following sequence of actions \"load object p0 into truck t0 at location l0-2, sail the ship t0 into city c0 from location l0-2 in city l0-0, remove the object p0 from the truck t0 and place it on the location l0-0, load the object p0 from location l0-0 onto the airplane a0, fly the airplane a0 from the airport l0-0 to the airport l1-0, remove the object p0 from the airplane a0 and place it on the location l1-0, load object p2 into truck t1 at location l1-2, navigate the truck t1 from its current location l1-2 in city c1 to the new location l1-0 within the same city place the object p0 into the truck t1 at location l1-0 remove the object p2 from the truck t1 and place it on the location l1-0 load the object p2 from location l1-0 onto the airplane a0 fly the airplane a0 from location l1-0 to location l2-0 fly airplane a0 from airport l2-0 to airport l0-0 unload the object p2 from the airplane a0 at location l0-0 place the object p2 into the truck t0 at location l0-0 navigate the truck t0 from its current location l0-0 in city c0 to the new location l0-1 within the same city offload the object p2 from the truck t0 at location l0-1 drive the truck t1 in city c1 from location l1-0 to location l1-2 offload the object p0 from the truck t1 at location l1-2 navigate the truck t2 from its current location l2-2 in city c2 to the new location l2-1 within the same city\" and the current state? A. The sequence is not valid. B. The sequence is applicable, but does not achieve the goal. C. The sequence is a plan. D. The sequence is not applicable.' + answer: "Let's think step by step. Step 1: For a sequence of actions to be a plan, all actions should be valid, applicable in sequence, and achieve the goal. Step 2: The action \"sail the ship t0 into city c0 from location l0-2 in city l0-0\" is not valid in this problem. **Final Answer**: A." diff --git a/lm-evaluation-harness/lm_eval/tasks/aexams/README.md b/lm-evaluation-harness/lm_eval/tasks/aexams/README.md new file mode 100644 index 0000000000000000000000000000000000000000..0c0461920c1cd51b6c3a4deb2af68843558116e1 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/aexams/README.md @@ -0,0 +1,53 @@ +# Arabic EXAMS + +### Paper + +EXAMS: a resource specialized in multilingual high school exam questions. +The original paper [EXAMS](https://aclanthology.org/2020.emnlp-main.438/) + +The Arabic EXAMS dataset includes five subjects + + - Islamic studies + - Biology + - Physics + - Science + - Social + +The original dataset [EXAMS-QA](https://github.com/mhardalov/exams-qa) + +EXAMS is a benchmark dataset for cross-lingual and multilingual question answering for high school examinations. +With 24,000 high-quality high school exam questions in 16 languages, covering 8 language families and 24 school subjects from Natural Sciences and Social Sciences, among others. +EXAMS offers unique fine-grained evaluation framework across multiple languages and subjects + +Homepage for Arabic EXAMS: [EXAMS Arabic Homepage](https://github.com/FreedomIntelligence/AceGPT/tree/main/eval/benchmark_eval/benchmarks/EXAMS_Arabic) + +### Citation + + +### Groups, Tags, and Tasks + +#### Groups + +- `aexams`: Arabic EXAMS dataset, including IslamicStudies, Biology, Science, Physics, Social subjects. + +#### Tasks + + +The following tasks evaluate subjects in Arabic EXAMS dataset using loglikelihood-based multiple-choice scoring: +- `aexams_IslamicStudies` +- `aexams_Biology` +- `aexams_Science` +- `aexams_Physics` +- `aexams_Social` + +### Checklist + +* [x] Is the task an existing benchmark in the literature? + * [x] Have you referenced the original paper that introduced the task? + * [x] If yes, does the original paper provide a reference implementation? + * [x] Yes, original implementation contributed by author of the benchmark + +If other tasks on this dataset are already supported: +* [x] Is the "Main" variant of this task clearly denoted? +* [x] Have you provided a short sentence in a README on what each new variant adds / evaluates? +* [x] Have you noted which, if any, published evaluation setups are matched by this variant? diff --git a/lm-evaluation-harness/lm_eval/tasks/aexams/_aexams.yaml b/lm-evaluation-harness/lm_eval/tasks/aexams/_aexams.yaml new file mode 100644 index 0000000000000000000000000000000000000000..59099b9c38c11e5391e031a2e07808a83d645938 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/aexams/_aexams.yaml @@ -0,0 +1,16 @@ +group: aexams +task: + - aexams_Biology + - aexams_IslamicStudies + - aexams_Physics + - aexams_Science + - aexams_Social +aggregate_metric_list: + - metric: acc + aggregation: mean + weight_by_size: true + - metric: acc_norm + aggregation: mean + weight_by_size: true +metadata: + version: 1.0 diff --git a/lm-evaluation-harness/lm_eval/tasks/aexams/_default_template_yaml b/lm-evaluation-harness/lm_eval/tasks/aexams/_default_template_yaml new file mode 100644 index 0000000000000000000000000000000000000000..3f7100ad70190a67bd86675ce7a15d88a5a5976a --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/aexams/_default_template_yaml @@ -0,0 +1,18 @@ +dataset_path: Hennara/aexams +test_split: test +fewshot_split: dev +fewshot_config: + sampler: first_n +output_type: multiple_choice +doc_to_text: "{{question.strip()}}\nA. {{A}}\nB. {{B}}\nC. {{C}}\nD. {{D}}\nالجواب：" +doc_to_choice: ["A", "B", "C", "D"] +doc_to_target: "{{['A', 'B', 'C', 'D'].index(answer)}}" +metric_list: + - metric: acc + aggregation: mean + higher_is_better: true + - metric: acc_norm + aggregation: mean + higher_is_better: true +metadata: + version: 1.0 diff --git a/lm-evaluation-harness/lm_eval/tasks/aexams/aexams_Biology.yaml b/lm-evaluation-harness/lm_eval/tasks/aexams/aexams_Biology.yaml new file mode 100644 index 0000000000000000000000000000000000000000..9ee2e33b5844ef438da4ac51bfd916af04cb53e6 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/aexams/aexams_Biology.yaml @@ -0,0 +1,4 @@ +"dataset_name": "Biology" +"description": "قم بالإجابة على مايلي في مجال العلوم الحيوية\n\n" +"include": "_default_template_yaml" +"task": "aexams_Biology" diff --git a/lm-evaluation-harness/lm_eval/tasks/aexams/aexams_IslamicStudies.yaml b/lm-evaluation-harness/lm_eval/tasks/aexams/aexams_IslamicStudies.yaml new file mode 100644 index 0000000000000000000000000000000000000000..831afc376ec25fdeddbb18bf5e4063d2e3c17ebf --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/aexams/aexams_IslamicStudies.yaml @@ -0,0 +1,4 @@ +"dataset_name": "IslamicStudies" +"description": "قم بالإجابة على مايلي في مجال العلوم الإسلامية \n\n" +"include": "_default_template_yaml" +"task": "aexams_IslamicStudies" diff --git a/lm-evaluation-harness/lm_eval/tasks/aexams/aexams_Physics.yaml b/lm-evaluation-harness/lm_eval/tasks/aexams/aexams_Physics.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f2764a06ef2680a1c81ccca0e76dcbcf1ba52672 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/aexams/aexams_Physics.yaml @@ -0,0 +1,4 @@ +"dataset_name": "Physics" +"description": "قم بالإجابة على مايلي في مجال الفيزياء \n\n" +"include": "_default_template_yaml" +"task": "aexams_Physics" diff --git a/lm-evaluation-harness/lm_eval/tasks/aexams/aexams_Science.yaml b/lm-evaluation-harness/lm_eval/tasks/aexams/aexams_Science.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c89dc8c8ca6d32b922483f48ee8da427e027a92b --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/aexams/aexams_Science.yaml @@ -0,0 +1,4 @@ +"dataset_name": "Science" +"description": "قم بالإجابة على مايلي في مجال العلوم \n\n" +"include": "_default_template_yaml" +"task": "aexams_Science" diff --git a/lm-evaluation-harness/lm_eval/tasks/aexams/aexams_Social.yaml b/lm-evaluation-harness/lm_eval/tasks/aexams/aexams_Social.yaml new file mode 100644 index 0000000000000000000000000000000000000000..3042a419e6e3902ddd0090028fc4b875a148a213 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/aexams/aexams_Social.yaml @@ -0,0 +1,4 @@ +"dataset_name": "Social" +"description": "قم بالإجابة على مايلي في مجال العلوم الإجتماعية \n\n" +"include": "_default_template_yaml" +"task": "aexams_Social" diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/README.md b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/README.md new file mode 100644 index 0000000000000000000000000000000000000000..cca14d968d2d87312d48fdb031e4a3518c9f915a --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/README.md @@ -0,0 +1,52 @@ +# MathQA + +### Paper + +IrokoBench: A New Benchmark for African Languages in the Age of Large Language Models +https://arxiv.org/pdf/2406.03368 + +IrokoBench is a human-translated benchmark dataset for 16 typologically diverse +low-resource African languages covering three tasks: natural language inference (AfriXNLI), +mathematical reasoning (AfriMGSM), and multi-choice knowledge-based QA (AfriMMLU). + + +### Citation + +``` +@misc{adelani2024irokobenchnewbenchmarkafrican, + title={IrokoBench: A New Benchmark for African Languages in the Age of Large Language Models}, + author={David Ifeoluwa Adelani and Jessica Ojo and Israel Abebe Azime and Jian Yun Zhuang and Jesujoba O. Alabi and Xuanli He and Millicent Ochieng and Sara Hooker and Andiswa Bukula and En-Shiun Annie Lee and Chiamaka Chukwuneke and Happy Buzaaba and Blessing Sibanda and Godson Kalipe and Jonathan Mukiibi and Salomon Kabongo and Foutse Yuehgoh and Mmasibidi Setaka and Lolwethu Ndolela and Nkiruka Odu and Rooweither Mabuya and Shamsuddeen Hassan Muhammad and Salomey Osei and Sokhar Samb and Tadesse Kebede Guge and Pontus Stenetorp}, + year={2024}, + eprint={2406.03368}, + archivePrefix={arXiv}, + primaryClass={cs.CL}, + url={https://arxiv.org/abs/2406.03368}, +} +``` + +### Groups and Tasks + +#### Groups + +* `afrimgsm`: All afrimgsm tasks +* `afrimgsm_direct`: afrimgsm_direct evaluates models performance on the curated dataset +* `afrimgsm_en_cot`: afrimgsm_en_cot includes 5-shot of exemplars for chain-of-thought approach +* `afrimgsm_translate`: afrimgsm_translate evaluates models in translate-test setting + +#### Tasks +* `afrimgsm_direct_{language_code}`: each task evaluates for one language +* `afrimgsm_en_cot_{language_code}`: each task evaluates for one language +* `afrimgsm_translate_{language_code}`: each task evaluates for one language + +### Checklist + +For adding novel benchmarks/datasets to the library: +* [x] Is the task an existing benchmark in the literature? + * [x] Have you referenced the original paper that introduced the task? + * [ ] If yes, does the original paper provide a reference implementation? If so, have you checked against the reference implementation and documented how to run such a test? + +If other tasks on this dataset are already supported: +* [x] Is the "Main" variant of this task clearly denoted? +* [x] Have you provided a short sentence in a README on what each new variant adds / evaluates? +* [x] Have you noted which, if any, published evaluation setups are matched by this variant? + * [x] Checked for equivalence with v0.3.0 LM Evaluation Harness diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/afrimgsm.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/afrimgsm.yaml new file mode 100644 index 0000000000000000000000000000000000000000..2316a748bd6f72f0c234544c016ddfd6b33fd9ff --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/afrimgsm.yaml @@ -0,0 +1,13 @@ +group: afrimgsm-irokobench +task: + - afrimgsm_tasks_prompt_1 + - afrimgsm_tasks_prompt_2 + - afrimgsm_tasks_prompt_3 + - afrimgsm_tasks_prompt_4 + - afrimgsm_tasks_prompt_5 +aggregate_metric_list: + - metric: acc + aggregation: mean + weight_by_size: true +metadata: + version: 2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_1/afrimgsm_amh.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_1/afrimgsm_amh.yaml new file mode 100644 index 0000000000000000000000000000000000000000..23007e3657c85b3b42ac5591180096c54740a240 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_1/afrimgsm_amh.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: amh +include: afrimgsm_yaml +task: afrimgsm_amh_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_1/afrimgsm_eng.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_1/afrimgsm_eng.yaml new file mode 100644 index 0000000000000000000000000000000000000000..1d5694225089b96dfeb06d331482bafa821cede8 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_1/afrimgsm_eng.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: eng +include: afrimgsm_yaml +task: afrimgsm_eng_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_1/afrimgsm_fra.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_1/afrimgsm_fra.yaml new file mode 100644 index 0000000000000000000000000000000000000000..04d57dbd329cedd535476cd7980ac5a201d84847 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_1/afrimgsm_fra.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: fra +include: afrimgsm_yaml +task: afrimgsm_fra_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_1/afrimgsm_swa.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_1/afrimgsm_swa.yaml new file mode 100644 index 0000000000000000000000000000000000000000..be6dba7151542a8b6ecb4e5cb1da18ab0d9121a3 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_1/afrimgsm_swa.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: swa +include: afrimgsm_yaml +task: afrimgsm_swa_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_1/afrimgsm_vai.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_1/afrimgsm_vai.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e04d28f0d33f1da1d6282431c8d4e1655a1f175b --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_1/afrimgsm_vai.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: vai +include: afrimgsm_yaml +task: afrimgsm_vai_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_ewe.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_ewe.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d0cd4926120ab7de80c13ba9b13bd327c81866bb --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_ewe.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: ewe +include: afrimgsm_yaml +task: afrimgsm_ewe_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_kin.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_kin.yaml new file mode 100644 index 0000000000000000000000000000000000000000..9c7b65251577ccb7bfa1fe74ddbe06114247dae5 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_kin.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: kin +include: afrimgsm_yaml +task: afrimgsm_kin_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_vai.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_vai.yaml new file mode 100644 index 0000000000000000000000000000000000000000..655b23dec64404bb271b558726ef5f279096f4c0 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_vai.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: vai +include: afrimgsm_yaml +task: afrimgsm_vai_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/gen_utils.py b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/gen_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..ecef389f3a4051e57b652f617b19ddd15d3c26ca --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/gen_utils.py @@ -0,0 +1,122 @@ +import argparse +import os + +import yaml + + +class FunctionTag: + def __init__(self, value): + self.value = value + + +def prompt_func(mode, lang): + prompt_map = { + "prompt_4": "Answer the given question with the step by step solution appropriate numerical value, ensuring that the response is " + "clear and without any supplementary information. \n\nQuestion: {{question}} \nStep by step answer: ", + "prompt_5": f"For mathematical questions provided in {lang} language. Supply the accurate step by step answer to the " + "provided question. \n\nQuestion: {{question}} \nStep by step answer: ", + } + return prompt_map[mode] + + +def gen_lang_yamls(output_dir: str, overwrite: bool, mode: str) -> None: + """ + Generate a yaml file for each language. + + :param output_dir: The directory to output the files to. + :param overwrite: Whether to overwrite files if they already exist. + """ + err = [] + languages = { + "eng": "English", + "amh": "Amharic", + "ibo": "Igbo", + "fra": "French", + "sna": "chiShona", + "wol": "Wolof", + "ewe": "Ewe", + "lin": "Lingala", + "lug": "Luganda", + "xho": "isiXhosa", + "kin": "Kinyarwanda", + "twi": "Twi", + "zul": "Zulu", + "orm": "Oromo", + "yor": "Yoruba", + "hau": "Hausa", + "sot": "Sesotho", + "swa": "Swahili", + "vai": "Vai", + } + + for lang in languages.keys(): + try: + file_name = f"afrimgsm_cot_{lang}.yaml" + task_name = f"afrimgsm_cot_{lang}_{mode}" + yaml_template = "afrimgsm_cot_yaml" + if "translate" in output_dir.split("/")[-1]: + file_name = f"afrimgsm_cot_translate_{lang}.yaml" + task_name = f"afrimgsm_cot_translate_{lang}_{mode}" + yaml_template = "afrimgsm_cot_translate_yaml" + if int(mode.split("_")[-1]) > 3: + yaml_details = { + "include": yaml_template, + "task": task_name, + "dataset_name": lang, + "doc_to_text": prompt_func(mode, languages[lang]), + } + else: + yaml_details = { + "include": yaml_template, + "task": task_name, + "dataset_name": lang, + } + os.makedirs(f"{output_dir}/{mode}", exist_ok=True) + with open( + f"{output_dir}/{mode}/{file_name}", + "w" if overwrite else "x", + encoding="utf8", + ) as f: + f.write("# Generated by utils.py\n") + yaml.dump( + yaml_details, + f, + allow_unicode=True, + ) + except FileExistsError: + err.append(file_name) + + if len(err) > 0: + raise FileExistsError( + "Files were not created because they already exist (use --overwrite flag):" + f" {', '.join(err)}" + ) + + +def main() -> None: + """Parse CLI args and generate language-specific yaml files.""" + parser = argparse.ArgumentParser() + parser.add_argument( + "--overwrite", + default=True, + action="store_true", + help="Overwrite files if they already exist", + ) + parser.add_argument( + "--output-dir", + default="./translate_cot", + help="Directory to write yaml files to", + ) + parser.add_argument( + "--mode", + default="prompt_5", + choices=["prompt_1", "prompt_2", "prompt_3", "prompt_4", "prompt_5"], + help="Prompt number", + ) + args = parser.parse_args() + + gen_lang_yamls(output_dir=args.output_dir, overwrite=args.overwrite, mode=args.mode) + + +if __name__ == "__main__": + main() diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/gen_yaml.sh b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/gen_yaml.sh new file mode 100644 index 0000000000000000000000000000000000000000..5c0132822a7f3ba68230762e0342838583c29bd9 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/gen_yaml.sh @@ -0,0 +1,7 @@ +#!/bin/bash + +# python utils.py --overwrite --output-dir direct --mode direct +# python utils.py --overwrite --output-dir direct_native --mode direct-native +# python utils.py --overwrite --output-dir en_cot --mode en-cot +# python utils.py --overwrite --output-dir native_cot --mode native-cot +python utils.py --overwrite --output-dir translate_direct --mode translate-direct diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/run.sh b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/run.sh new file mode 100644 index 0000000000000000000000000000000000000000..075500be33775dc49288ce7f7180604c7c6f99ce --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/run.sh @@ -0,0 +1,6 @@ +lm_eval --model hf \ + --model_args pretrained="google/gemma-7b" --tasks afrimgsm_en_cot_eng,mgsm_en_cot_en,afrimgsm_native_cot_eng,mgsm_native_cot_en,afrimgsm_direct_eng,mgsm_direct_en,afrimgsm_direct_native_eng \ + --device cuda:0 \ + --batch_size 1 \ + --verbosity DEBUG \ + --limit 5 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/utils.py b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..0dd336f8b3cbeb85d5854beb923578f215f91632 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/utils.py @@ -0,0 +1,219 @@ +import argparse + +import yaml + + +languages = [ + "eng", + "amh", + "ibo", + "fra", + "sna", + "lin", + "wol", + "ewe", + "lug", + "xho", + "kin", + "twi", + "zul", + "orm", + "yor", + "hau", + "sot", + "swa", +] + +languages_REGEX = { + "eng": "The answer is (\\-?[0-9\\.\\,]+)", + "amh": "መልሱ (\\-?[0-9\\.\\,]+)", + "ibo": "Azịza ya bụ (\\-?[0-9\\.\\,]+)", + "fra": "La réponse est(\\-?[0-9\\.\\,]+)", + "sna": "Mhinduro kumubvunzo ndi (\\-?[0-9\\.\\,]+)", + "lin": "Eyano ezali (\\-?[0-9\\.\\,]+)", + "wol": "Tontu li (\\-?[0-9\\.\\,]+)", + "ewe": "ŋuɖoɖoae nye (\\-?[0-9\\.\\,]+)", + "lug": "Ansa eri (\\-?[0-9\\.\\,]+)", + "xho": "Impendulo ngu (\\-?[0-9\\.\\,]+)", + "kin": "Igisubizo ni (\\-?[0-9\\.\\,]+)", + "twi": "Ne nnyiano yɛ (\\-?[0-9\\.\\,]+)", + "zul": "Impendulo ithi (\\-?[0-9\\.\\,]+)", + "orm": "Deebiin isaa (\\-?[0-9\\.\\,]+)", + "yor": "Ìdáhùn náà ni (\\-?[0-9\\.\\,]+)", + "hau": "Amsar ita ce (\\-?[0-9\\.\\,]+)", + "sot": "Karabo ke (\\-?[0-9\\.\\,]+)", + "swa": "Jibu ni (\\-?[0-9\\.\\,]+)", +} + +LANGUAGES = {} + +for lang in languages: + if lang == "amh": + LANGUAGES[lang] = { # English + "QUESTION": "ጥያቄ:", + "ANSWER": "በቅደም ተከተል መልስ:", + "DIRECT": "Answer:", + "REGEX": languages_REGEX[lang], + } + elif lang == "yor": + LANGUAGES[lang] = { # English + "QUESTION": "Ìbéèrè:", + "ANSWER": "Ìdáhùn lẹ́sẹsẹ:", + "DIRECT": "Answer:", + "REGEX": languages_REGEX[lang], + } + + else: + LANGUAGES[lang] = { # English + "QUESTION": "Question:", + "ANSWER": "Step-by-Step Answer:", + "DIRECT": "Answer:", + "REGEX": languages_REGEX[lang], + } + + +def add_regex_pattern(regex_pattern): + if regex_pattern is None: + return {} + return { + "filter_list": [ + { + "name": "strict-match", + "filter": [ + { + "function": "regex", + "regex_pattern": f"""{regex_pattern}""", + }, + { + "function": "take_first", + }, + ], + }, + { + "name": "flexible-extract", + "filter": [ + { + "function": "regex", + "regex_pattern": """(-?[$0-9.,]{2,})|(-?[0-9]+)""", + "group_select": -1, + }, + { + "function": "take_first", + }, + ], + }, + ], + } + + +def gen_lang_yamls(output_dir: str, overwrite: bool, mode: str) -> None: + """ + Generate a yaml file for each language. + + :param output_dir: The directory to output the files to. + :param overwrite: Whether to overwrite files if they already exist. + """ + err = [] + for lang in LANGUAGES.keys(): + try: + yaml_template = "cot_yaml" + filter_list = {} + DELIMITER = None + if mode == "direct": + ANSWER = LANGUAGES["eng"]["DIRECT"] + QUESTION = LANGUAGES["eng"]["QUESTION"] + REGEX = None + task_name = f"afrimgsm_direct_{lang}" + yaml_template = "direct_yaml" + if mode == "direct-native": + ANSWER = LANGUAGES[lang]["DIRECT"] + QUESTION = LANGUAGES[lang]["QUESTION"] + REGEX = None + task_name = f"afrimgsm_direct_native_{lang}" + yaml_template = "direct_native_yaml" + elif mode == "native-cot": + ANSWER = LANGUAGES[lang]["ANSWER"] + REGEX = LANGUAGES[lang]["REGEX"] + QUESTION = LANGUAGES[lang]["QUESTION"] + task_name = f"afrimgsm_native_cot_{lang}" + filter_list = add_regex_pattern(REGEX) + DELIMITER = "" if lang in ["zh", "ja"] else None + elif mode == "en-cot": + ANSWER = LANGUAGES["eng"]["ANSWER"] + REGEX = LANGUAGES["eng"]["REGEX"] + QUESTION = LANGUAGES["eng"]["QUESTION"] + task_name = f"afrimgsm_en_cot_{lang}" + elif mode == "translate-direct": + ANSWER = LANGUAGES["eng"]["DIRECT"] + QUESTION = LANGUAGES["eng"]["QUESTION"] + REGEX = None + task_name = f"afrimgsm_translate_direct_{lang}" + yaml_template = "translate_direct_yaml" + + file_name = f"{task_name}.yaml" + ANSWER_TO_SKIP = len(LANGUAGES[lang]["ANSWER"]) + 1 + with open( + f"{output_dir}/{file_name}", "w" if overwrite else "x", encoding="utf8" + ) as f: + f.write("# Generated by utils.py\n") + yaml.dump( + { + "include": yaml_template, + "dataset_name": lang, + "task": f"{task_name}", + "doc_to_text": f"""{{% if answer is not none %}}""" + f"""{{{{question+"\\n{ANSWER}"}}}}""" + f"""{{% else %}}""" + f"""{{{{"{QUESTION} "+question+"\\n{ANSWER}"}}}}""" + f"""{{% endif %}}""", + "doc_to_target": f"""{{% if answer is not none %}}""" + f"""{{{{answer[{ANSWER_TO_SKIP}:]}}}}""" + f"""{{% else %}}""" + f"""{{{{answer_number|string}}}}""" + f"""{{% endif %}}""", + **filter_list, + "generation_kwargs": { + "until": [QUESTION, "", "<|im_end|>"], + "do_sample": False, + }, + **({"target_delimiter": DELIMITER} if DELIMITER else {}), + }, + f, + allow_unicode=True, + width=float("inf"), + ) + except FileExistsError: + err.append(file_name) + + if len(err) > 0: + raise FileExistsError( + "Files were not created because they already exist (use --overwrite flag):" + f" {', '.join(err)}" + ) + + +def main() -> None: + """Parse CLI args and generate language-specific yaml files.""" + parser = argparse.ArgumentParser() + parser.add_argument( + "--overwrite", + default=False, + action="store_true", + help="Overwrite files if they already exist", + ) + parser.add_argument( + "--output-dir", default=".", help="Directory to write yaml files to" + ) + parser.add_argument( + "--mode", + default="native-cot", + choices=["direct", "direct-native", "native-cot", "en-cot", "translate-direct"], + help="Mode of chain-of-thought", + ) + args = parser.parse_args() + + gen_lang_yamls(output_dir=args.output_dir, overwrite=args.overwrite, mode=args.mode) + + +if __name__ == "__main__": + main() diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimmlu/translate/prompt_5/afrimmlu_translate_wol.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimmlu/translate/prompt_5/afrimmlu_translate_wol.yaml new file mode 100644 index 0000000000000000000000000000000000000000..2cd495e81df06383612c1278b18faeb0ac5c567f --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimmlu/translate/prompt_5/afrimmlu_translate_wol.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: wol +include: afrimmlu_translate +task: afrimmlu_translate_wol_prompt_5 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimmlu/translate/prompt_5/utils.py b/lm-evaluation-harness/lm_eval/tasks/afrimmlu/translate/prompt_5/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..147225bb70653d67663ac1762a7cd6246c4e9f22 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimmlu/translate/prompt_5/utils.py @@ -0,0 +1,28 @@ +from lm_eval.utils import weighted_f1_score + + +def doc_to_choice(doc): + choices = eval(doc["choices"]) + return choices + + +def doc_to_text(doc): + output = """Given your proficiency in {subject}, please answer the subsequent multiple-choice question with 'A', 'B', 'C', or 'D'. +Question: {question} +Choices: + A: {choice1} + B: {choice2} + C: {choice3} + D: {choice4} +Answer: """ + + choices = eval(doc["choices"]) + text = output.format( + subject=doc["subject"], + question=doc["question"], + choice1=choices[0], + choice2=choices[1], + choice3=choices[2], + choice4=choices[3], + ) + return text diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/en-direct/afrixnli_en_direct_swa.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/en-direct/afrixnli_en_direct_swa.yaml new file mode 100644 index 0000000000000000000000000000000000000000..bd81dde5c2de1b7cd199ae1975253a6edce73f62 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/en-direct/afrixnli_en_direct_swa.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: swa +include: afrixnli_en_direct_yaml +task: afrixnli_en_direct_swa diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/en-direct/afrixnli_en_direct_twi.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/en-direct/afrixnli_en_direct_twi.yaml new file mode 100644 index 0000000000000000000000000000000000000000..49dd2b1e0a9915e8c1773603af640bc6a33e89a7 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/en-direct/afrixnli_en_direct_twi.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: twi +include: afrixnli_en_direct_yaml +task: afrixnli_en_direct_twi diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/en-direct/afrixnli_en_direct_wol.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/en-direct/afrixnli_en_direct_wol.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0de5a2fedc41c60981d0b77e7d9d67c57855de08 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/en-direct/afrixnli_en_direct_wol.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: wol +include: afrixnli_en_direct_yaml +task: afrixnli_en_direct_wol diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/en-direct/afrixnli_en_direct_xho.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/en-direct/afrixnli_en_direct_xho.yaml new file mode 100644 index 0000000000000000000000000000000000000000..38e4ca57a413e0cd0b830f72adba13ffa281367b --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/en-direct/afrixnli_en_direct_xho.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: xho +include: afrixnli_en_direct_yaml +task: afrixnli_en_direct_xho diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/en-direct/afrixnli_en_direct_yor.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/en-direct/afrixnli_en_direct_yor.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c43ffac0c320911ea5d76e3c5c39a331d489a19f --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/en-direct/afrixnli_en_direct_yor.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: yor +include: afrixnli_en_direct_yaml +task: afrixnli_en_direct_yor diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/en-direct/afrixnli_en_direct_zul.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/en-direct/afrixnli_en_direct_zul.yaml new file mode 100644 index 0000000000000000000000000000000000000000..80d078937bc19bda8bbbb3dd6e8c425fac3e9b23 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/en-direct/afrixnli_en_direct_zul.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: zul +include: afrixnli_en_direct_yaml +task: afrixnli_en_direct_zul diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/en-direct/utils.py b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/en-direct/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..5d1ac19e19b2e855c957e75f1c778366dfbc7e55 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/en-direct/utils.py @@ -0,0 +1,6 @@ +from lm_eval.utils import weighted_f1_score + + +def doc_to_target(doc): + replacements = {0: "True", 1: "Neither", 2: "False"} + return replacements[doc["label"]] diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/afrixnli_native_direct_amh.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/afrixnli_native_direct_amh.yaml new file mode 100644 index 0000000000000000000000000000000000000000..87b517906c0129b4aaabc518766e89dd6f70f505 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/afrixnli_native_direct_amh.yaml @@ -0,0 +1,6 @@ +# Generated by utils.py +dataset_name: amh +doc_to_choice: '{{[premise+", ትክክል? አዎ, "+hypothesis,premise+", ትክክል? እንዲሁም, "+hypothesis,premise+", + ትክክል? አይ, "+hypothesis]}}' +include: afrixnli_native_direct_yaml +task: afrixnli_native_direct_amh diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/afrixnli_native_direct_fra.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/afrixnli_native_direct_fra.yaml new file mode 100644 index 0000000000000000000000000000000000000000..1e6e32cc165ce73e99ae5480debe0183dbb2351a --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/afrixnli_native_direct_fra.yaml @@ -0,0 +1,6 @@ +# Generated by utils.py +dataset_name: fra +doc_to_choice: '{{[premise+", correct? Oui, "+hypothesis,premise+", correct? Aussi, + "+hypothesis,premise+", correct? Non, "+hypothesis]}}' +include: afrixnli_native_direct_yaml +task: afrixnli_native_direct_fra diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/afrixnli_native_direct_hau.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/afrixnli_native_direct_hau.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b243a5de37f970dc92f27112280332cd2c5256cd --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/afrixnli_native_direct_hau.yaml @@ -0,0 +1,6 @@ +# Generated by utils.py +dataset_name: hau +doc_to_choice: '{{[premise+", Daidai? Ee, "+hypothesis,premise+", Daidai? Haka kuma, + "+hypothesis,premise+", Daidai? A''a, "+hypothesis]}}' +include: afrixnli_native_direct_yaml +task: afrixnli_native_direct_hau diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/afrixnli_native_direct_kin.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/afrixnli_native_direct_kin.yaml new file mode 100644 index 0000000000000000000000000000000000000000..3333c12019585f6d089635cb7d7ab5eb9ad906d6 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/afrixnli_native_direct_kin.yaml @@ -0,0 +1,6 @@ +# Generated by utils.py +dataset_name: kin +doc_to_choice: '{{[premise+", Nibyo? Yego, "+hypothesis,premise+", Nibyo? Na none, + "+hypothesis,premise+", Nibyo? Oya, "+hypothesis]}}' +include: afrixnli_native_direct_yaml +task: afrixnli_native_direct_kin diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/afrixnli_native_direct_lin.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/afrixnli_native_direct_lin.yaml new file mode 100644 index 0000000000000000000000000000000000000000..95060d6869cb8e6a3340f9d54bc5257244666f3c --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/afrixnli_native_direct_lin.yaml @@ -0,0 +1,6 @@ +# Generated by utils.py +dataset_name: lin +doc_to_choice: '{{[premise+", Malamu? Iyo, "+hypothesis,premise+", Malamu? Lisusu, + "+hypothesis,premise+", Malamu? Te, "+hypothesis]}}' +include: afrixnli_native_direct_yaml +task: afrixnli_native_direct_lin diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/afrixnli_native_direct_lug.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/afrixnli_native_direct_lug.yaml new file mode 100644 index 0000000000000000000000000000000000000000..97b6d00ec8b4ed0d9e1c9aabe133cc0b70141dbb --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/afrixnli_native_direct_lug.yaml @@ -0,0 +1,6 @@ +# Generated by utils.py +dataset_name: lug +doc_to_choice: '{{[premise+", Kituufu? Yee, "+hypothesis,premise+", Kituufu? N’ekirala, + "+hypothesis,premise+", Kituufu? Nedda, "+hypothesis]}}' +include: afrixnli_native_direct_yaml +task: afrixnli_native_direct_lug diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/afrixnli_native_direct_orm.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/afrixnli_native_direct_orm.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f9c25496da9cc81a3c82c8ae2a83621bf839e56a --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/afrixnli_native_direct_orm.yaml @@ -0,0 +1,6 @@ +# Generated by utils.py +dataset_name: orm +doc_to_choice: '{{[premise+", Sirrii? Eeyyee, "+hypothesis,premise+", Sirrii? Akkasumas, + "+hypothesis,premise+", Sirrii? Lakki, "+hypothesis]}}' +include: afrixnli_native_direct_yaml +task: afrixnli_native_direct_orm diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/afrixnli_native_direct_sna.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/afrixnli_native_direct_sna.yaml new file mode 100644 index 0000000000000000000000000000000000000000..be2b2617ccdec63258607be20a6db2d958f018b1 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/afrixnli_native_direct_sna.yaml @@ -0,0 +1,6 @@ +# Generated by utils.py +dataset_name: sna +doc_to_choice: '{{[premise+", Chokwadi? Hongu, "+hypothesis,premise+", Chokwadi? Uye, + "+hypothesis,premise+", Chokwadi? Kwete, "+hypothesis]}}' +include: afrixnli_native_direct_yaml +task: afrixnli_native_direct_sna diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/afrixnli_native_direct_sot.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/afrixnli_native_direct_sot.yaml new file mode 100644 index 0000000000000000000000000000000000000000..092961e0f8e39bb94a152aae00651b9ae49eebfb --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/afrixnli_native_direct_sot.yaml @@ -0,0 +1,6 @@ +# Generated by utils.py +dataset_name: sot +doc_to_choice: '{{[premise+", Nepile? E, "+hypothesis,premise+", Nepile? Hape, "+hypothesis,premise+", + Nepile? Tjhe, "+hypothesis]}}' +include: afrixnli_native_direct_yaml +task: afrixnli_native_direct_sot diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/afrixnli_native_direct_swa.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/afrixnli_native_direct_swa.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c8b1e2afa2c1b267c803565caa9cc13dc8d8f506 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/afrixnli_native_direct_swa.yaml @@ -0,0 +1,6 @@ +# Generated by utils.py +dataset_name: swa +doc_to_choice: '{{[premise+", Sahihi? Ndiyo, "+hypothesis,premise+", Sahihi? Pia, + "+hypothesis,premise+", Sahihi? Hapana, "+hypothesis]}}' +include: afrixnli_native_direct_yaml +task: afrixnli_native_direct_swa diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/afrixnli_native_direct_twi.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/afrixnli_native_direct_twi.yaml new file mode 100644 index 0000000000000000000000000000000000000000..4d3141d63a84b5a93002bd416583eb444543c031 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/afrixnli_native_direct_twi.yaml @@ -0,0 +1,6 @@ +# Generated by utils.py +dataset_name: twi +doc_to_choice: '{{[premise+", Nifa? Aane, "+hypothesis,premise+", Nifa? Anaasɛ, "+hypothesis,premise+", + Nifa? Daabi, "+hypothesis]}}' +include: afrixnli_native_direct_yaml +task: afrixnli_native_direct_twi diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/afrixnli_native_direct_wol.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/afrixnli_native_direct_wol.yaml new file mode 100644 index 0000000000000000000000000000000000000000..1239fa47086826050a23d493c3e7069327a0e516 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/afrixnli_native_direct_wol.yaml @@ -0,0 +1,6 @@ +# Generated by utils.py +dataset_name: wol +doc_to_choice: '{{[premise+", Dëgg? Waaw, "+hypothesis,premise+", Dëgg? Itam, "+hypothesis,premise+", + Dëgg? Déet, "+hypothesis]}}' +include: afrixnli_native_direct_yaml +task: afrixnli_native_direct_wol diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/afrixnli_native_direct_xho.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/afrixnli_native_direct_xho.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f6f91f6e079d1138e374c7094bd76ef4743ec5b4 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/afrixnli_native_direct_xho.yaml @@ -0,0 +1,6 @@ +# Generated by utils.py +dataset_name: xho +doc_to_choice: '{{[premise+", Ichanekile? Ewe, "+hypothesis,premise+", Ichanekile? + Kananjalo, "+hypothesis,premise+", Ichanekile? Hayi, "+hypothesis]}}' +include: afrixnli_native_direct_yaml +task: afrixnli_native_direct_xho diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/afrixnli_native_direct_yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/afrixnli_native_direct_yaml new file mode 100644 index 0000000000000000000000000000000000000000..d5ec109bbd64b5cdbbd27329fa0bd7c67767cf5c --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/afrixnli_native_direct_yaml @@ -0,0 +1,25 @@ +tag: + - afrixnli + - afrixnli_native_direct +dataset_path: masakhane/afrixnli +dataset_name: null +output_type: multiple_choice +validation_split: validation +test_split: test +fewshot_split: validation +doc_to_target: label +doc_to_text: "" +metric_list: + - metric: f1 + aggregation: !function utils.weighted_f1_score + average: weighted + higher_is_better: True + ignore_case: true + ignore_punctuation: true + - metric: acc + aggregation: mean + higher_is_better: true + ignore_case: true + ignore_punctuation: true +metadata: + version: 1.0 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/afrixnli_native_direct_yor.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/afrixnli_native_direct_yor.yaml new file mode 100644 index 0000000000000000000000000000000000000000..2648bf57bce8ffa5d28578094b477c6b8b166446 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/afrixnli_native_direct_yor.yaml @@ -0,0 +1,6 @@ +# Generated by utils.py +dataset_name: yor +doc_to_choice: '{{[premise+", Òótọ́? Bẹ́ẹ̀ni, "+hypothesis,premise+", Òótọ́? Àti pé, + "+hypothesis,premise+", Òótọ́? Rárá, "+hypothesis]}}' +include: afrixnli_native_direct_yaml +task: afrixnli_native_direct_yor diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/afrixnli_native_direct_zul.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/afrixnli_native_direct_zul.yaml new file mode 100644 index 0000000000000000000000000000000000000000..48261c60b28fa4c157b511153b609c840fea80e8 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/afrixnli_native_direct_zul.yaml @@ -0,0 +1,6 @@ +# Generated by utils.py +dataset_name: zul +doc_to_choice: '{{[premise+", Kulungile? Yebo, "+hypothesis,premise+", Kulungile? + Futhi, "+hypothesis,premise+", Kulungile? Cha, "+hypothesis]}}' +include: afrixnli_native_direct_yaml +task: afrixnli_native_direct_zul diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/utils.py b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..3e735e2deb1f9c53152c072615aebe8ba3acb90b --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/native-direct/utils.py @@ -0,0 +1 @@ +from lm_eval.utils import weighted_f1_score diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_amh.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_amh.yaml new file mode 100644 index 0000000000000000000000000000000000000000..94fb2bdcb6f44646e6711dfaa38d7d0f66c767f5 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_amh.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: amh +include: afrixnli_translate_yaml +task: afrixnli_translate_amh diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_ewe.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_ewe.yaml new file mode 100644 index 0000000000000000000000000000000000000000..55d5b470a2fdc47c73ac9ebeabbc6bdf388db0f2 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_ewe.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: ewe +include: afrixnli_translate_yaml +task: afrixnli_translate_ewe diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_fra.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_fra.yaml new file mode 100644 index 0000000000000000000000000000000000000000..bd5903357dbd029bbc5a3d88c47e75ab05b4da41 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_fra.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: fra +include: afrixnli_translate_yaml +task: afrixnli_translate_fra diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_hau.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_hau.yaml new file mode 100644 index 0000000000000000000000000000000000000000..ddc7a375e03210ad02090a0279fe767e67d76c8e --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_hau.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: hau +include: afrixnli_translate_yaml +task: afrixnli_translate_hau diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_ibo.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_ibo.yaml new file mode 100644 index 0000000000000000000000000000000000000000..2487f15a4a75ede35ab29300b6764fabd325e139 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_ibo.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: ibo +include: afrixnli_translate_yaml +task: afrixnli_translate_ibo diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_kin.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_kin.yaml new file mode 100644 index 0000000000000000000000000000000000000000..ebae340f5bf3b21a5d72c1ed4f6bad6223834d27 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_kin.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: kin +include: afrixnli_translate_yaml +task: afrixnli_translate_kin diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_lin.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_lin.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0ad2ea078f78d509c452850ec1fdeef2c1f96325 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_lin.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: lin +include: afrixnli_translate_yaml +task: afrixnli_translate_lin diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_lug.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_lug.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c9ab91826d3f8b64370b061b58dcd5cd1b5d0da8 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_lug.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: lug +include: afrixnli_translate_yaml +task: afrixnli_translate_lug diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_orm.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_orm.yaml new file mode 100644 index 0000000000000000000000000000000000000000..879228120a74794e894cad0b6d32ccb0b35ad473 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_orm.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: orm +include: afrixnli_translate_yaml +task: afrixnli_translate_orm diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_sna.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_sna.yaml new file mode 100644 index 0000000000000000000000000000000000000000..69756c268c21a7228ed87cfc41522b5f2f549bf1 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_sna.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: sna +include: afrixnli_translate_yaml +task: afrixnli_translate_sna diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_sot.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_sot.yaml new file mode 100644 index 0000000000000000000000000000000000000000..64b5cb29c770a1380a69001edf0026f47a0509a7 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_sot.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: sot +include: afrixnli_translate_yaml +task: afrixnli_translate_sot diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_swa.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_swa.yaml new file mode 100644 index 0000000000000000000000000000000000000000..ea6307131bfe14bdf9951b929556b0e911bed25f --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_swa.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: swa +include: afrixnli_translate_yaml +task: afrixnli_translate_swa diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_twi.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_twi.yaml new file mode 100644 index 0000000000000000000000000000000000000000..5cfd32e21ffd7208af52822be3bbeacd1676efab --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_twi.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: twi +include: afrixnli_translate_yaml +task: afrixnli_translate_twi diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_wol.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_wol.yaml new file mode 100644 index 0000000000000000000000000000000000000000..be1188e5cc7c941099bf25d9d7b71eba768dcf9a --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_wol.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: wol +include: afrixnli_translate_yaml +task: afrixnli_translate_wol diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_xho.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_xho.yaml new file mode 100644 index 0000000000000000000000000000000000000000..428ff3bbd2dccc0f60bb3818860d8426f9f70739 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_xho.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: xho +include: afrixnli_translate_yaml +task: afrixnli_translate_xho diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_yaml new file mode 100644 index 0000000000000000000000000000000000000000..3f1df47cf06db0a549f279ee81ea5e8d8945a85e --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_yaml @@ -0,0 +1,32 @@ +tag: + - afrixnli + - afrixnli_translate +dataset_path: masakhane/afrixnli-translate-test +dataset_name: null +output_type: multiple_choice +test_split: test +doc_to_text: "{{premise}}\nQuestion: {{hypothesis}} True, False, or Neither?\nAnswer:" +# True = entailment +# False = contradiction +# Neither = neutral +doc_to_target: !function utils.doc_to_target +doc_to_choice: + - "True" + - "Neither" + - "False" +should_decontaminate: true +doc_to_decontamination_query: premise +metric_list: + - metric: f1 + aggregation: !function utils.weighted_f1_score + average: weighted + higher_is_better: True + ignore_case: true + ignore_punctuation: true + - metric: acc + aggregation: mean + higher_is_better: true + ignore_case: true + ignore_punctuation: true +metadata: + version: 1.0 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_yor.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_yor.yaml new file mode 100644 index 0000000000000000000000000000000000000000..4f07a41a5b1b1f48ff85110aa3c6d1197a51f437 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_yor.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: yor +include: afrixnli_translate_yaml +task: afrixnli_translate_yor diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_zul.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_zul.yaml new file mode 100644 index 0000000000000000000000000000000000000000..7a57632bcafc9da38097eea7fbad89c14fbd12e9 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/afrixnli_translate_zul.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: zul +include: afrixnli_translate_yaml +task: afrixnli_translate_zul diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/utils.py b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..5d1ac19e19b2e855c957e75f1c778366dfbc7e55 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/anli prompt/translate/utils.py @@ -0,0 +1,6 @@ +from lm_eval.utils import weighted_f1_score + + +def doc_to_target(doc): + replacements = {0: "True", 1: "Neither", 2: "False"} + return replacements[doc["label"]] diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/afrixnli.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/afrixnli.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d85ccd128f752f7a1ab566aa28e90d5bbf545b66 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/afrixnli.yaml @@ -0,0 +1,13 @@ +group: afrixnli-irokobench +task: + - afrixnli_tasks_prompt_1 + - afrixnli_tasks_prompt_2 + - afrixnli_tasks_prompt_3 + - afrixnli_tasks_prompt_4 + - afrixnli_tasks_prompt_5 +aggregate_metric_list: + - metric: acc + aggregation: mean + weight_by_size: true +metadata: + version: 2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_amh.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_amh.yaml new file mode 100644 index 0000000000000000000000000000000000000000..39f727b4ccb7c07eb0b2f6b8d2472764446767d4 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_amh.yaml @@ -0,0 +1,15 @@ +# Generated by utils.py +dataset_name: amh +doc_to_text: 'Please identify whether the premise entails or contradicts the hypothesis + in the following premise and hypothesis. The answer should be exact entailment, + contradiction, or neutral. + + + Premise: {premise} + + Hypothesis: {hypothesis} + + + Is it entailment, contradiction, or neutral?' +include: afrixnli_yaml +task: afrixnli_amh_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_eng.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_eng.yaml new file mode 100644 index 0000000000000000000000000000000000000000..593c57a34ec0f01d3c03e447acda48cd1644231b --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_eng.yaml @@ -0,0 +1,15 @@ +# Generated by utils.py +dataset_name: eng +doc_to_text: 'Please identify whether the premise entails or contradicts the hypothesis + in the following premise and hypothesis. The answer should be exact entailment, + contradiction, or neutral. + + + Premise: {premise} + + Hypothesis: {hypothesis} + + + Is it entailment, contradiction, or neutral?' +include: afrixnli_yaml +task: afrixnli_eng_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_ewe.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_ewe.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b6a10baae753575ebb228a6df34e4faf364efea1 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_ewe.yaml @@ -0,0 +1,15 @@ +# Generated by utils.py +dataset_name: ewe +doc_to_text: 'Please identify whether the premise entails or contradicts the hypothesis + in the following premise and hypothesis. The answer should be exact entailment, + contradiction, or neutral. + + + Premise: {premise} + + Hypothesis: {hypothesis} + + + Is it entailment, contradiction, or neutral?' +include: afrixnli_yaml +task: afrixnli_ewe_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_fra.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_fra.yaml new file mode 100644 index 0000000000000000000000000000000000000000..08b2b5243633f276487d8d5595382211870eedf9 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_fra.yaml @@ -0,0 +1,15 @@ +# Generated by utils.py +dataset_name: fra +doc_to_text: 'Please identify whether the premise entails or contradicts the hypothesis + in the following premise and hypothesis. The answer should be exact entailment, + contradiction, or neutral. + + + Premise: {premise} + + Hypothesis: {hypothesis} + + + Is it entailment, contradiction, or neutral?' +include: afrixnli_yaml +task: afrixnli_fra_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_hau.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_hau.yaml new file mode 100644 index 0000000000000000000000000000000000000000..fe234b72694fdfde8474863d81265e851a350368 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_hau.yaml @@ -0,0 +1,15 @@ +# Generated by utils.py +dataset_name: hau +doc_to_text: 'Please identify whether the premise entails or contradicts the hypothesis + in the following premise and hypothesis. The answer should be exact entailment, + contradiction, or neutral. + + + Premise: {premise} + + Hypothesis: {hypothesis} + + + Is it entailment, contradiction, or neutral?' +include: afrixnli_yaml +task: afrixnli_hau_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_ibo.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_ibo.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d282e0e5f84433b77f8954acaa4e65c9ccbf5ba4 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_ibo.yaml @@ -0,0 +1,15 @@ +# Generated by utils.py +dataset_name: ibo +doc_to_text: 'Please identify whether the premise entails or contradicts the hypothesis + in the following premise and hypothesis. The answer should be exact entailment, + contradiction, or neutral. + + + Premise: {premise} + + Hypothesis: {hypothesis} + + + Is it entailment, contradiction, or neutral?' +include: afrixnli_yaml +task: afrixnli_ibo_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_kin.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_kin.yaml new file mode 100644 index 0000000000000000000000000000000000000000..cfdff6c8c64e6a91a869386f71b6f6024c0ac156 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_kin.yaml @@ -0,0 +1,15 @@ +# Generated by utils.py +dataset_name: kin +doc_to_text: 'Please identify whether the premise entails or contradicts the hypothesis + in the following premise and hypothesis. The answer should be exact entailment, + contradiction, or neutral. + + + Premise: {premise} + + Hypothesis: {hypothesis} + + + Is it entailment, contradiction, or neutral?' +include: afrixnli_yaml +task: afrixnli_kin_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_lin.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_lin.yaml new file mode 100644 index 0000000000000000000000000000000000000000..410cb29f80366d78d0ec4fb9e240a9df0ec20372 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_lin.yaml @@ -0,0 +1,15 @@ +# Generated by utils.py +dataset_name: lin +doc_to_text: 'Please identify whether the premise entails or contradicts the hypothesis + in the following premise and hypothesis. The answer should be exact entailment, + contradiction, or neutral. + + + Premise: {premise} + + Hypothesis: {hypothesis} + + + Is it entailment, contradiction, or neutral?' +include: afrixnli_yaml +task: afrixnli_lin_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_lug.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_lug.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b5665e37cce68d072cf8b64be5a787aed23fd70b --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_lug.yaml @@ -0,0 +1,15 @@ +# Generated by utils.py +dataset_name: lug +doc_to_text: 'Please identify whether the premise entails or contradicts the hypothesis + in the following premise and hypothesis. The answer should be exact entailment, + contradiction, or neutral. + + + Premise: {premise} + + Hypothesis: {hypothesis} + + + Is it entailment, contradiction, or neutral?' +include: afrixnli_yaml +task: afrixnli_lug_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_orm.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_orm.yaml new file mode 100644 index 0000000000000000000000000000000000000000..12751c7f93de40a8bc91431de573be9f99868ab4 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_orm.yaml @@ -0,0 +1,15 @@ +# Generated by utils.py +dataset_name: orm +doc_to_text: 'Please identify whether the premise entails or contradicts the hypothesis + in the following premise and hypothesis. The answer should be exact entailment, + contradiction, or neutral. + + + Premise: {premise} + + Hypothesis: {hypothesis} + + + Is it entailment, contradiction, or neutral?' +include: afrixnli_yaml +task: afrixnli_orm_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_sna.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_sna.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d00bbb6f9d146effdf5de3112db4c56f79002166 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_sna.yaml @@ -0,0 +1,15 @@ +# Generated by utils.py +dataset_name: sna +doc_to_text: 'Please identify whether the premise entails or contradicts the hypothesis + in the following premise and hypothesis. The answer should be exact entailment, + contradiction, or neutral. + + + Premise: {premise} + + Hypothesis: {hypothesis} + + + Is it entailment, contradiction, or neutral?' +include: afrixnli_yaml +task: afrixnli_sna_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_sot.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_sot.yaml new file mode 100644 index 0000000000000000000000000000000000000000..2ae346aed8ff8e53b94252385443b54fe4364595 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_sot.yaml @@ -0,0 +1,15 @@ +# Generated by utils.py +dataset_name: sot +doc_to_text: 'Please identify whether the premise entails or contradicts the hypothesis + in the following premise and hypothesis. The answer should be exact entailment, + contradiction, or neutral. + + + Premise: {premise} + + Hypothesis: {hypothesis} + + + Is it entailment, contradiction, or neutral?' +include: afrixnli_yaml +task: afrixnli_sot_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_swa.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_swa.yaml new file mode 100644 index 0000000000000000000000000000000000000000..ca6729bf80cad6e95027c7c0e994cd1da14d0d6d --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_swa.yaml @@ -0,0 +1,15 @@ +# Generated by utils.py +dataset_name: swa +doc_to_text: 'Please identify whether the premise entails or contradicts the hypothesis + in the following premise and hypothesis. The answer should be exact entailment, + contradiction, or neutral. + + + Premise: {premise} + + Hypothesis: {hypothesis} + + + Is it entailment, contradiction, or neutral?' +include: afrixnli_yaml +task: afrixnli_swa_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_twi.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_twi.yaml new file mode 100644 index 0000000000000000000000000000000000000000..7dc85428dab2ed5da0cb6fa17b0a428088f346f1 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_twi.yaml @@ -0,0 +1,15 @@ +# Generated by utils.py +dataset_name: twi +doc_to_text: 'Please identify whether the premise entails or contradicts the hypothesis + in the following premise and hypothesis. The answer should be exact entailment, + contradiction, or neutral. + + + Premise: {premise} + + Hypothesis: {hypothesis} + + + Is it entailment, contradiction, or neutral?' +include: afrixnli_yaml +task: afrixnli_twi_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_wol.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_wol.yaml new file mode 100644 index 0000000000000000000000000000000000000000..78ef254aeed1d393efca6390fa574aa41eac5f21 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_wol.yaml @@ -0,0 +1,15 @@ +# Generated by utils.py +dataset_name: wol +doc_to_text: 'Please identify whether the premise entails or contradicts the hypothesis + in the following premise and hypothesis. The answer should be exact entailment, + contradiction, or neutral. + + + Premise: {premise} + + Hypothesis: {hypothesis} + + + Is it entailment, contradiction, or neutral?' +include: afrixnli_yaml +task: afrixnli_wol_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_xho.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_xho.yaml new file mode 100644 index 0000000000000000000000000000000000000000..cb0a8527741f4b8f3ac1fb7c2741a6cf5e2c64ae --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_xho.yaml @@ -0,0 +1,15 @@ +# Generated by utils.py +dataset_name: xho +doc_to_text: 'Please identify whether the premise entails or contradicts the hypothesis + in the following premise and hypothesis. The answer should be exact entailment, + contradiction, or neutral. + + + Premise: {premise} + + Hypothesis: {hypothesis} + + + Is it entailment, contradiction, or neutral?' +include: afrixnli_yaml +task: afrixnli_xho_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_yaml new file mode 100644 index 0000000000000000000000000000000000000000..81c9eeaa5af0740cc32122519f671c4d0425c080 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_yaml @@ -0,0 +1,30 @@ +tag: + - afrixnli_tasks + - afrixnli_tasks_prompt_1 +dataset_path: masakhane/afrixnli +dataset_name: null +output_type: multiple_choice +validation_split: validation +test_split: test +fewshot_split: validation +doc_to_target: !function utils.doc_to_target +doc_to_choice: + - "entailment" + - "neutral" + - "contradiction" +should_decontaminate: true +doc_to_decontamination_query: premise +metric_list: + - metric: f1 + aggregation: !function utils.weighted_f1_score + average: weighted + higher_is_better: True + ignore_case: true + ignore_punctuation: true + - metric: acc + aggregation: mean + higher_is_better: true + ignore_case: true + ignore_punctuation: true +metadata: + version: 1.0 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_yor.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_yor.yaml new file mode 100644 index 0000000000000000000000000000000000000000..473aea37a7b036d7ef219eca756482cd2bac754b --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_yor.yaml @@ -0,0 +1,15 @@ +# Generated by utils.py +dataset_name: yor +doc_to_text: 'Please identify whether the premise entails or contradicts the hypothesis + in the following premise and hypothesis. The answer should be exact entailment, + contradiction, or neutral. + + + Premise: {premise} + + Hypothesis: {hypothesis} + + + Is it entailment, contradiction, or neutral?' +include: afrixnli_yaml +task: afrixnli_yor_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_zul.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_zul.yaml new file mode 100644 index 0000000000000000000000000000000000000000..fa07a8c991d56a0cef3dc8453017649952715f8a --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/afrixnli_zul.yaml @@ -0,0 +1,15 @@ +# Generated by utils.py +dataset_name: zul +doc_to_text: 'Please identify whether the premise entails or contradicts the hypothesis + in the following premise and hypothesis. The answer should be exact entailment, + contradiction, or neutral. + + + Premise: {premise} + + Hypothesis: {hypothesis} + + + Is it entailment, contradiction, or neutral?' +include: afrixnli_yaml +task: afrixnli_zul_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/utils.py b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..d97a0a288508e817ab695e637fb157a08c813808 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_1/utils.py @@ -0,0 +1,19 @@ +from lm_eval.utils import weighted_f1_score + + +def doc_to_text(doc): + output = """Please identify whether the premise entails or contradicts the hypothesis in the following premise + and hypothesis. The answer should be exact entailment, contradiction, or neutral. + + Premise: {premise} + Hypothesis: {hypothesis} + + Is it entailment, contradiction, or neutral?""" + + text = output.format(premise=doc["premise"], hypothesis=doc["hypothesis"]) + return text + + +def doc_to_target(doc): + replacements = {0: "entailment", 1: "neutral", 2: "contradiction"} + return replacements[doc["label"]] diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_amh.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_amh.yaml new file mode 100644 index 0000000000000000000000000000000000000000..fbf916b25d43db2fdc476c27fe5f2e8e02c45625 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_amh.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: amh +include: afrixnli_yaml +task: afrixnli_amh_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_eng.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_eng.yaml new file mode 100644 index 0000000000000000000000000000000000000000..dfa8ebfe8815a58c1a043b328a22f762a739b9d2 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_eng.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: eng +include: afrixnli_yaml +task: afrixnli_eng_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_ewe.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_ewe.yaml new file mode 100644 index 0000000000000000000000000000000000000000..995ef3e65894548266a72f45417222e2760e30fa --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_ewe.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: ewe +include: afrixnli_yaml +task: afrixnli_ewe_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_fra.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_fra.yaml new file mode 100644 index 0000000000000000000000000000000000000000..ce72588c19901c04ee82479206f54816fa358915 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_fra.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: fra +include: afrixnli_yaml +task: afrixnli_fra_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_hau.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_hau.yaml new file mode 100644 index 0000000000000000000000000000000000000000..369ee58bedfd98fd95c63510c3a84eec10238df0 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_hau.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: hau +include: afrixnli_yaml +task: afrixnli_hau_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_ibo.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_ibo.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e118c613ebf2d0388298fe6ba750923816ba4af6 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_ibo.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: ibo +include: afrixnli_yaml +task: afrixnli_ibo_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_kin.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_kin.yaml new file mode 100644 index 0000000000000000000000000000000000000000..81f6d803d6762f5a6b86dae00ec0b26040a943ac --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_kin.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: kin +include: afrixnli_yaml +task: afrixnli_kin_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_lin.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_lin.yaml new file mode 100644 index 0000000000000000000000000000000000000000..2d99c2eb57aedcce988f37415c414882d5bb4186 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_lin.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: lin +include: afrixnli_yaml +task: afrixnli_lin_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_lug.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_lug.yaml new file mode 100644 index 0000000000000000000000000000000000000000..31325539e1dd778da5e057436aeb3b60d7531a58 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_lug.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: lug +include: afrixnli_yaml +task: afrixnli_lug_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_orm.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_orm.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c4ad555afafe2b99470d706e0eb46dc8256037fb --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_orm.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: orm +include: afrixnli_yaml +task: afrixnli_orm_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_sna.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_sna.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a780b0c428c822c08d5bb16dd909cb883da494a2 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_sna.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: sna +include: afrixnli_yaml +task: afrixnli_sna_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_sot.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_sot.yaml new file mode 100644 index 0000000000000000000000000000000000000000..94e78880d31b48ce8dc4e562a9d9cc3643208535 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_sot.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: sot +include: afrixnli_yaml +task: afrixnli_sot_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_twi.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_twi.yaml new file mode 100644 index 0000000000000000000000000000000000000000..4219b81ee8a1de24535cf2cd6eae4643e660d0de --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_twi.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: twi +include: afrixnli_yaml +task: afrixnli_twi_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_wol.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_wol.yaml new file mode 100644 index 0000000000000000000000000000000000000000..546b17904959bd83c1f26617a9a13b79fc654a55 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_wol.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: wol +include: afrixnli_yaml +task: afrixnli_wol_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_xho.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_xho.yaml new file mode 100644 index 0000000000000000000000000000000000000000..649c61df93eef6b6828d7da9c5a672bc5fce9611 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_xho.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: xho +include: afrixnli_yaml +task: afrixnli_xho_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_yaml new file mode 100644 index 0000000000000000000000000000000000000000..cfab642bf9175d3066879680618e95f097a609a2 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_yaml @@ -0,0 +1,34 @@ +tag: + - afrixnli_tasks + - afrixnli_tasks_prompt_2 +dataset_path: masakhane/afrixnli +dataset_name: null +output_type: multiple_choice +validation_split: validation +test_split: test +fewshot_split: validation +doc_to_text: "{{premise}}\nQuestion: {{hypothesis}} True, False, or Neither?\nAnswer:" +# True = entailment +# False = contradiction +# Neither = neutral +doc_to_target: !function utils.doc_to_target +doc_to_choice: + - "True" + - "Neither" + - "False" +should_decontaminate: true +doc_to_decontamination_query: premise +metric_list: + - metric: f1 + aggregation: !function utils.weighted_f1_score + average: weighted + higher_is_better: True + ignore_case: true + ignore_punctuation: true + - metric: acc + aggregation: mean + higher_is_better: true + ignore_case: true + ignore_punctuation: true +metadata: + version: 1.0 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_yor.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_yor.yaml new file mode 100644 index 0000000000000000000000000000000000000000..53f23ace6bbb25c6457f7bd4e5b760b7ebb8b298 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_yor.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: yor +include: afrixnli_yaml +task: afrixnli_yor_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_zul.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_zul.yaml new file mode 100644 index 0000000000000000000000000000000000000000..dd89fe131e26f4b0d1dedb1b86aeac72f8f706d6 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/afrixnli_zul.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: zul +include: afrixnli_yaml +task: afrixnli_zul_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/utils.py b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..5d1ac19e19b2e855c957e75f1c778366dfbc7e55 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_2/utils.py @@ -0,0 +1,6 @@ +from lm_eval.utils import weighted_f1_score + + +def doc_to_target(doc): + replacements = {0: "True", 1: "Neither", 2: "False"} + return replacements[doc["label"]] diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_3/afrixnli_amh.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_3/afrixnli_amh.yaml new file mode 100644 index 0000000000000000000000000000000000000000..3ff9f99c187a914fde7514c7c4caf49cb63c4186 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_3/afrixnli_amh.yaml @@ -0,0 +1,8 @@ +# Generated by utils.py +dataset_name: amh +doc_to_text: "Given the following premise and hypothesis in Amharic, identify if the\ + \ premise entails, contradicts, or is neutral towards the hypothesis. Please respond\ + \ with exact 'entailment', 'contradiction', or 'neutral'. \n\nPremise: {{premise}}\ + \ \nHypothesis: {{hypothesis}}" +include: afrixnli_yaml +task: afrixnli_amh_prompt_3 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_3/afrixnli_eng.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_3/afrixnli_eng.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a53aea6dbd9c3fedd9a812fc8f698b5f16d41bf3 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_3/afrixnli_eng.yaml @@ -0,0 +1,8 @@ +# Generated by utils.py +dataset_name: eng +doc_to_text: "Given the following premise and hypothesis in English, identify if the\ + \ premise entails, contradicts, or is neutral towards the hypothesis. Please respond\ + \ with exact 'entailment', 'contradiction', or 'neutral'. \n\nPremise: {{premise}}\ + \ \nHypothesis: {{hypothesis}}" +include: afrixnli_yaml +task: afrixnli_eng_prompt_3 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_3/afrixnli_ewe.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_3/afrixnli_ewe.yaml new file mode 100644 index 0000000000000000000000000000000000000000..54b58ae6e972774be55be9988a20d6962a7e56ff --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_3/afrixnli_ewe.yaml @@ -0,0 +1,8 @@ +# Generated by utils.py +dataset_name: ewe +doc_to_text: "Given the following premise and hypothesis in Ewe, identify if the premise\ + \ entails, contradicts, or is neutral towards the hypothesis. Please respond with\ + \ exact 'entailment', 'contradiction', or 'neutral'. \n\nPremise: {{premise}} \n\ + Hypothesis: {{hypothesis}}" +include: afrixnli_yaml +task: afrixnli_ewe_prompt_3 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_3/afrixnli_fra.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_3/afrixnli_fra.yaml new file mode 100644 index 0000000000000000000000000000000000000000..fedb519ec8421d77aedb24215de643544614bf70 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_3/afrixnli_fra.yaml @@ -0,0 +1,8 @@ +# Generated by utils.py +dataset_name: fra +doc_to_text: "Given the following premise and hypothesis in French, identify if the\ + \ premise entails, contradicts, or is neutral towards the hypothesis. Please respond\ + \ with exact 'entailment', 'contradiction', or 'neutral'. \n\nPremise: {{premise}}\ + \ \nHypothesis: {{hypothesis}}" +include: afrixnli_yaml +task: afrixnli_fra_prompt_3 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_3/afrixnli_hau.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_3/afrixnli_hau.yaml new file mode 100644 index 0000000000000000000000000000000000000000..3a9ebb95181426ab8a9138267a8400c37825e76e --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_3/afrixnli_hau.yaml @@ -0,0 +1,8 @@ +# Generated by utils.py +dataset_name: hau +doc_to_text: "Given the following premise and hypothesis in Hausa, identify if the\ + \ premise entails, contradicts, or is neutral towards the hypothesis. Please respond\ + \ with exact 'entailment', 'contradiction', or 'neutral'. \n\nPremise: {{premise}}\ + \ \nHypothesis: {{hypothesis}}" +include: afrixnli_yaml +task: afrixnli_hau_prompt_3 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_3/afrixnli_ibo.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_3/afrixnli_ibo.yaml new file mode 100644 index 0000000000000000000000000000000000000000..6b61f7678a3ab596bbda1b6039129e2e71b6bda6 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_3/afrixnli_ibo.yaml @@ -0,0 +1,8 @@ +# Generated by utils.py +dataset_name: ibo +doc_to_text: "Given the following premise and hypothesis in Igbo, identify if the\ + \ premise entails, contradicts, or is neutral towards the hypothesis. Please respond\ + \ with exact 'entailment', 'contradiction', or 'neutral'. \n\nPremise: {{premise}}\ + \ \nHypothesis: {{hypothesis}}" +include: afrixnli_yaml +task: afrixnli_ibo_prompt_3 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_3/afrixnli_lin.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_3/afrixnli_lin.yaml new file mode 100644 index 0000000000000000000000000000000000000000..697c439fa18a1ecd80e41fca14bc0836d956cfde --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_3/afrixnli_lin.yaml @@ -0,0 +1,8 @@ +# Generated by utils.py +dataset_name: lin +doc_to_text: "Given the following premise and hypothesis in Lingala, identify if the\ + \ premise entails, contradicts, or is neutral towards the hypothesis. Please respond\ + \ with exact 'entailment', 'contradiction', or 'neutral'. \n\nPremise: {{premise}}\ + \ \nHypothesis: {{hypothesis}}" +include: afrixnli_yaml +task: afrixnli_lin_prompt_3 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_3/afrixnli_lug.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_3/afrixnli_lug.yaml new file mode 100644 index 0000000000000000000000000000000000000000..1b5667c0720473f0ab7703b17777b5ced0459381 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_3/afrixnli_lug.yaml @@ -0,0 +1,8 @@ +# Generated by utils.py +dataset_name: lug +doc_to_text: "Given the following premise and hypothesis in Luganda, identify if the\ + \ premise entails, contradicts, or is neutral towards the hypothesis. Please respond\ + \ with exact 'entailment', 'contradiction', or 'neutral'. \n\nPremise: {{premise}}\ + \ \nHypothesis: {{hypothesis}}" +include: afrixnli_yaml +task: afrixnli_lug_prompt_3 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_3/afrixnli_orm.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_3/afrixnli_orm.yaml new file mode 100644 index 0000000000000000000000000000000000000000..37a6d843e51870f6ad2845e1f385b6aacef2fab2 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_3/afrixnli_orm.yaml @@ -0,0 +1,8 @@ +# Generated by utils.py +dataset_name: orm +doc_to_text: "Given the following premise and hypothesis in Oromo, identify if the\ + \ premise entails, contradicts, or is neutral towards the hypothesis. Please respond\ + \ with exact 'entailment', 'contradiction', or 'neutral'. \n\nPremise: {{premise}}\ + \ \nHypothesis: {{hypothesis}}" +include: afrixnli_yaml +task: afrixnli_orm_prompt_3 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_3/afrixnli_sna.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_3/afrixnli_sna.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c7e0f0b05000c40fbebca19921a2486a57d8054a --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_3/afrixnli_sna.yaml @@ -0,0 +1,8 @@ +# Generated by utils.py +dataset_name: sna +doc_to_text: "Given the following premise and hypothesis in chiShona, identify if\ + \ the premise entails, contradicts, or is neutral towards the hypothesis. Please\ + \ respond with exact 'entailment', 'contradiction', or 'neutral'. \n\nPremise: {{premise}}\ + \ \nHypothesis: {{hypothesis}}" +include: afrixnli_yaml +task: afrixnli_sna_prompt_3 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_3/afrixnli_sot.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_3/afrixnli_sot.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0c0ccd9e64ee0b807fef742c873126666327f276 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_3/afrixnli_sot.yaml @@ -0,0 +1,8 @@ +# Generated by utils.py +dataset_name: sot +doc_to_text: "Given the following premise and hypothesis in Sesotho, identify if the\ + \ premise entails, contradicts, or is neutral towards the hypothesis. Please respond\ + \ with exact 'entailment', 'contradiction', or 'neutral'. \n\nPremise: {{premise}}\ + \ \nHypothesis: {{hypothesis}}" +include: afrixnli_yaml +task: afrixnli_sot_prompt_3 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_3/afrixnli_swa.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_3/afrixnli_swa.yaml new file mode 100644 index 0000000000000000000000000000000000000000..dabd96ef2e1cbf6df83432cc057382d40e448ff7 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_3/afrixnli_swa.yaml @@ -0,0 +1,8 @@ +# Generated by utils.py +dataset_name: swa +doc_to_text: "Given the following premise and hypothesis in Swahili, identify if the\ + \ premise entails, contradicts, or is neutral towards the hypothesis. Please respond\ + \ with exact 'entailment', 'contradiction', or 'neutral'. \n\nPremise: {{premise}}\ + \ \nHypothesis: {{hypothesis}}" +include: afrixnli_yaml +task: afrixnli_swa_prompt_3 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_4/afrixnli_amh.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_4/afrixnli_amh.yaml new file mode 100644 index 0000000000000000000000000000000000000000..63b05465144af310263939fea2b8335672dbb7ae --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_4/afrixnli_amh.yaml @@ -0,0 +1,9 @@ +# Generated by utils.py +dataset_name: amh +doc_to_text: "You are an expert in Natural Language Inference (NLI) specializing in\ + \ the Amharic language.\nAnalyze the premise and hypothesis given in Amharic, and\ + \ determine the relationship between them.\n Respond with one of the following options:\ + \ 'entailment', 'contradiction', or 'neutral'. \n\nPremise: {{premise}} \nHypothesis:\ + \ {{hypothesis}}" +include: afrixnli_yaml +task: afrixnli_amh_prompt_4 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_4/afrixnli_eng.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_4/afrixnli_eng.yaml new file mode 100644 index 0000000000000000000000000000000000000000..1ecb06d10497274dde56ab73302525add553254a --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_4/afrixnli_eng.yaml @@ -0,0 +1,9 @@ +# Generated by utils.py +dataset_name: eng +doc_to_text: "You are an expert in Natural Language Inference (NLI) specializing in\ + \ the English language.\nAnalyze the premise and hypothesis given in English, and\ + \ determine the relationship between them.\n Respond with one of the following options:\ + \ 'entailment', 'contradiction', or 'neutral'. \n\nPremise: {{premise}} \nHypothesis:\ + \ {{hypothesis}}" +include: afrixnli_yaml +task: afrixnli_eng_prompt_4 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_4/afrixnli_hau.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_4/afrixnli_hau.yaml new file mode 100644 index 0000000000000000000000000000000000000000..811a0fca1364f55e5ba3dfe37e7d9c99e7090e6a --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/direct/prompt_4/afrixnli_hau.yaml @@ -0,0 +1,8 @@ +# Generated by utils.py +dataset_name: hau +doc_to_text: "You are an expert in Natural Language Inference (NLI) specializing in\ + \ the Hausa language.\nAnalyze the premise and hypothesis given in Hausa, and determine\ + \ the relationship between them.\n Respond with one of the following options: 'entailment',\ + \ 'contradiction', or 'neutral'. \n\nPremise: {{premise}} \nHypothesis: {{hypothesis}}" +include: afrixnli_yaml +task: afrixnli_hau_prompt_4 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrixnli/translate/prompt_5/afrixnli_translate_yor.yaml b/lm-evaluation-harness/lm_eval/tasks/afrixnli/translate/prompt_5/afrixnli_translate_yor.yaml new file mode 100644 index 0000000000000000000000000000000000000000..107c663428d39e3eaa565315a40e4aa5f4b53201 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrixnli/translate/prompt_5/afrixnli_translate_yor.yaml @@ -0,0 +1,6 @@ +# Generated by utils.py +dataset_name: yor +doc_to_text: "Based on the given statement, is the following claim 'true', 'false',\ + \ or 'inconclusive'. \nStatement: {{premise}} \nClaim: {{hypothesis}}" +include: afrixnli_translate_yaml +task: afrixnli_translate_yor_prompt_5 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/adr/prompt_3/afridiacritics_bbj.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/adr/prompt_3/afridiacritics_bbj.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a50b40c535d778cb4bd564455fbfdcf43415a53d --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/adr/prompt_3/afridiacritics_bbj.yaml @@ -0,0 +1,6 @@ +# Generated by utils.py +dataset_name: bbj +doc_to_text: 'This text is in Gbomala. Restore all diacritical marks to their proper + places in the following sentence: {{text}}. Return output sentence only' +include: afridiacritics_yaml +task: afridiacritics_bbj_prompt_3 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/adr/prompt_3/afridiacritics_fon.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/adr/prompt_3/afridiacritics_fon.yaml new file mode 100644 index 0000000000000000000000000000000000000000..5b0909ce9dd69f46cdca75ebdc325f452b25a462 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/adr/prompt_3/afridiacritics_fon.yaml @@ -0,0 +1,6 @@ +# Generated by utils.py +dataset_name: fon +doc_to_text: 'This text is in Fon. Restore all diacritical marks to their proper places + in the following sentence: {{text}}. Return output sentence only' +include: afridiacritics_yaml +task: afridiacritics_fon_prompt_3 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/adr/prompt_3/afridiacritics_ibo.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/adr/prompt_3/afridiacritics_ibo.yaml new file mode 100644 index 0000000000000000000000000000000000000000..04d1df0e1f07ac7c082bd75b1ce93959e0e0d56d --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/adr/prompt_3/afridiacritics_ibo.yaml @@ -0,0 +1,6 @@ +# Generated by utils.py +dataset_name: ibo +doc_to_text: 'This text is in Igbo. Restore all diacritical marks to their proper + places in the following sentence: {{text}}. Return output sentence only' +include: afridiacritics_yaml +task: afridiacritics_ibo_prompt_3 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/adr/prompt_3/afridiacritics_wol.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/adr/prompt_3/afridiacritics_wol.yaml new file mode 100644 index 0000000000000000000000000000000000000000..576e0845188b523be7ec2f342a440174aa496263 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/adr/prompt_3/afridiacritics_wol.yaml @@ -0,0 +1,6 @@ +# Generated by utils.py +dataset_name: wol +doc_to_text: 'This text is in Wolof. Restore all diacritical marks to their proper + places in the following sentence: {{text}}. Return output sentence only' +include: afridiacritics_yaml +task: afridiacritics_wol_prompt_3 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/adr/prompt_3/afridiacritics_yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/adr/prompt_3/afridiacritics_yaml new file mode 100644 index 0000000000000000000000000000000000000000..0a27eeef2d37880527c7b99f1fa9296f843b72a0 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/adr/prompt_3/afridiacritics_yaml @@ -0,0 +1,25 @@ +tag: +- adr_tasks +- adr_prompt_3 +dataset_path: masakhane/diacritics-restoration +dataset_kwargs: {trust_remote_code: True} +doc_to_target: target +output_type: generate_until +fewshot_split: dev +test_split: test +training_split: train +metric_list: + - metric: bleu + aggregation: bleu + higher_is_better: true + - metric: chrf + aggregation: chrf + higher_is_better: true +generation_kwargs: + do_sample: false + until: + - '' + - + - <|im_end|> +metadata: + version: 1.0 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/adr/prompt_3/afridiacritics_yor.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/adr/prompt_3/afridiacritics_yor.yaml new file mode 100644 index 0000000000000000000000000000000000000000..169c110872d6fdc2d2d41b6472fe30d93934f5df --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/adr/prompt_3/afridiacritics_yor.yaml @@ -0,0 +1,6 @@ +# Generated by utils.py +dataset_name: yor +doc_to_text: 'This text is in Yoruba. Restore all diacritical marks to their proper + places in the following sentence: {{text}}. Return output sentence only' +include: afridiacritics_yaml +task: afridiacritics_yor_prompt_3 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/adr/prompt_4/afridiacritics_ibo.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/adr/prompt_4/afridiacritics_ibo.yaml new file mode 100644 index 0000000000000000000000000000000000000000..367e387ae7456f57d604b1fe3ac032084b16fb98 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/adr/prompt_4/afridiacritics_ibo.yaml @@ -0,0 +1,6 @@ +# Generated by utils.py +dataset_name: ibo +doc_to_text: 'You are a linguist specializing in diacritical marks for Igbo. Add the + appropriate diacritics to this Igbo sentence: {{text}}. Return output sentence only' +include: afridiacritics_yaml +task: afridiacritics_ibo_prompt_4 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/adr/prompt_4/afridiacritics_wol.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/adr/prompt_4/afridiacritics_wol.yaml new file mode 100644 index 0000000000000000000000000000000000000000..23fb81e754445e8745d5d67720707af7d502e3df --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/adr/prompt_4/afridiacritics_wol.yaml @@ -0,0 +1,7 @@ +# Generated by utils.py +dataset_name: wol +doc_to_text: 'You are a linguist specializing in diacritical marks for Wolof. Add + the appropriate diacritics to this Wolof sentence: {{text}}. Return output sentence + only' +include: afridiacritics_yaml +task: afridiacritics_wol_prompt_4 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/adr/prompt_4/afridiacritics_yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/adr/prompt_4/afridiacritics_yaml new file mode 100644 index 0000000000000000000000000000000000000000..6ae62e9d3384d3ee1bff044dbfd1cb23275ae517 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/adr/prompt_4/afridiacritics_yaml @@ -0,0 +1,25 @@ +tag: +- adr_tasks +- adr_prompt_4 +dataset_path: masakhane/diacritics-restoration +dataset_kwargs: {trust_remote_code: True} +doc_to_target: target +output_type: generate_until +fewshot_split: dev +test_split: test +training_split: train +metric_list: + - metric: bleu + aggregation: bleu + higher_is_better: true + - metric: chrf + aggregation: chrf + higher_is_better: true +generation_kwargs: + do_sample: false + until: + - '' + - + - <|im_end|> +metadata: + version: 1.0 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/adr/prompt_4/afridiacritics_yor.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/adr/prompt_4/afridiacritics_yor.yaml new file mode 100644 index 0000000000000000000000000000000000000000..21e3a53fefcb4ae41eb00a406a3319f14ed60aba --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/adr/prompt_4/afridiacritics_yor.yaml @@ -0,0 +1,7 @@ +# Generated by utils.py +dataset_name: yor +doc_to_text: 'You are a linguist specializing in diacritical marks for Yoruba. Add + the appropriate diacritics to this Yoruba sentence: {{text}}. Return output sentence + only' +include: afridiacritics_yaml +task: afridiacritics_yor_prompt_4 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/adr/prompt_5/afridiacritics_bbj.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/adr/prompt_5/afridiacritics_bbj.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b1bcc833c73d0a789700c1b50b8636163620ed27 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/adr/prompt_5/afridiacritics_bbj.yaml @@ -0,0 +1,8 @@ +# Generated by utils.py +dataset_name: bbj +doc_to_text: 'You are a linguist specializing in diacritical marks for Gbomala. Diacritics + are essential for proper pronunciation and meaning in Gbomala. You are tasked with + converting Gbomala sentences without diacritics into their correctly accented forms. + Here''s the input: {{text}}. Return output sentence only' +include: afridiacritics_yaml +task: afridiacritics_bbj_prompt_5 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/adr/prompt_5/afridiacritics_fon.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/adr/prompt_5/afridiacritics_fon.yaml new file mode 100644 index 0000000000000000000000000000000000000000..3a1c55f813b4c4b7d08daff74cf32040b85e2b35 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/adr/prompt_5/afridiacritics_fon.yaml @@ -0,0 +1,8 @@ +# Generated by utils.py +dataset_name: fon +doc_to_text: 'You are a linguist specializing in diacritical marks for Fon. Diacritics + are essential for proper pronunciation and meaning in Fon. You are tasked with converting + Fon sentences without diacritics into their correctly accented forms. Here''s the + input: {{text}}. Return output sentence only' +include: afridiacritics_yaml +task: afridiacritics_fon_prompt_5 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/adr/prompt_5/afridiacritics_ibo.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/adr/prompt_5/afridiacritics_ibo.yaml new file mode 100644 index 0000000000000000000000000000000000000000..6cc9865dca7a2d8889df88d73abfb54b615089f4 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/adr/prompt_5/afridiacritics_ibo.yaml @@ -0,0 +1,8 @@ +# Generated by utils.py +dataset_name: ibo +doc_to_text: 'You are a linguist specializing in diacritical marks for Igbo. Diacritics + are essential for proper pronunciation and meaning in Igbo. You are tasked with + converting Igbo sentences without diacritics into their correctly accented forms. + Here''s the input: {{text}}. Return output sentence only' +include: afridiacritics_yaml +task: afridiacritics_ibo_prompt_5 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/adr/prompt_5/afridiacritics_wol.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/adr/prompt_5/afridiacritics_wol.yaml new file mode 100644 index 0000000000000000000000000000000000000000..fed10a7031ac71a948720b15eff1677df411934c --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/adr/prompt_5/afridiacritics_wol.yaml @@ -0,0 +1,8 @@ +# Generated by utils.py +dataset_name: wol +doc_to_text: 'You are a linguist specializing in diacritical marks for Wolof. Diacritics + are essential for proper pronunciation and meaning in Wolof. You are tasked with + converting Wolof sentences without diacritics into their correctly accented forms. + Here''s the input: {{text}}. Return output sentence only' +include: afridiacritics_yaml +task: afridiacritics_wol_prompt_5 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/adr/prompt_5/afridiacritics_yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/adr/prompt_5/afridiacritics_yaml new file mode 100644 index 0000000000000000000000000000000000000000..aaad3306e7270e78cdd2f83dd8ffeb790520134d --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/adr/prompt_5/afridiacritics_yaml @@ -0,0 +1,25 @@ +tag: +- adr_tasks +- adr_prompt_5 +dataset_path: masakhane/diacritics-restoration +dataset_kwargs: {trust_remote_code: True} +doc_to_target: target +output_type: generate_until +fewshot_split: dev +test_split: test +training_split: train +metric_list: + - metric: bleu + aggregation: bleu + higher_is_better: true + - metric: chrf + aggregation: chrf + higher_is_better: true +generation_kwargs: + do_sample: false + until: + - '' + - + - <|im_end|> +metadata: + version: 1.0 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/adr/prompt_5/afridiacritics_yor.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/adr/prompt_5/afridiacritics_yor.yaml new file mode 100644 index 0000000000000000000000000000000000000000..bd1c9007a394de95a19c1a09b398614366537a1f --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/adr/prompt_5/afridiacritics_yor.yaml @@ -0,0 +1,8 @@ +# Generated by utils.py +dataset_name: yor +doc_to_text: 'You are a linguist specializing in diacritical marks for Yoruba. Diacritics + are essential for proper pronunciation and meaning in Yoruba. You are tasked with + converting Yoruba sentences without diacritics into their correctly accented forms. + Here''s the input: {{text}}. Return output sentence only' +include: afridiacritics_yaml +task: afridiacritics_yor_prompt_5 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/README.md b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/README.md new file mode 100644 index 0000000000000000000000000000000000000000..8730d7c8d8d68b6b83dfad3d4f584534b048d111 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/README.md @@ -0,0 +1,24 @@ +# + +## Paper +Title: `AfriQA: Cross-lingual Open-Retrieval Question Answering for African Languages` + +Paper Link: https://arxiv.org/abs/2305.06897 + +## Abstract +>AfriQA is the first cross-lingual question answering (QA) dataset with a focus on African languages. The dataset includes over 12,000 XOR QA examples across 10 African languages, making it an invaluable resource for developing more equitable QA technology. African languages have historically been underserved in the digital landscape, with far less in-language content available online. This makes it difficult for QA systems to provide accurate information to users in their native language. However, cross-lingual open-retrieval question answering (XOR QA) systems can help fill this gap by retrieving answer content from other languages. AfriQA focuses specifically on African languages where cross-lingual answer content is the only high-coverage source of information. Previous datasets have primarily focused on languages where cross-lingual QA augments coverage from the target language, but AfriQA highlights the importance of African languages as a realistic use case for XOR QA. + +HomePage: https://github.com/masakhane-io/afriqa + +### Citation + +``` +@misc{ogundepo2023afriqa, + title={AfriQA: Cross-lingual Open-Retrieval Question Answering for African Languages}, + author={Odunayo Ogundepo and Tajuddeen R. Gwadabe and Clara E. Rivera and Jonathan H. Clark and Sebastian Ruder and David Ifeoluwa Adelani and Bonaventure F. P. Dossou and Abdou Aziz DIOP and Claytone Sikasote and Gilles Hacheme and Happy Buzaaba and Ignatius Ezeani and Rooweither Mabuya and Salomey Osei and Chris Emezue and Albert Njoroge Kahira and Shamsuddeen H. Muhammad and Akintunde Oladipo and Abraham Toluwase Owodunni and Atnafu Lambebo Tonja and Iyanuoluwa Shode and Akari Asai and Tunde Oluwaseyi Ajayi and Clemencia Siro and Steven Arthur and Mofetoluwa Adeyemi and Orevaoghene Ahia and Aremu Anuoluwapo and Oyinkansola Awosan and Chiamaka Chukwuneke and Bernard Opoku and Awokoya Ayodele and Verrah Otiende and Christine Mwase and Boyd Sinkala and Andre Niyongabo Rubungo and Daniel A. Ajisafe and Emeka Felix Onwuegbuzia and Habib Mbow and Emile Niyomutabazi and Eunice Mukonde and Falalu Ibrahim Lawan and Ibrahim Said Ahmad and Jesujoba O. Alabi and Martin Namukombo and Mbonu Chinedu and Mofya Phiri and Neo Putini and Ndumiso Mngoma and Priscilla A. Amuok and Ruqayya Nasir Iro and Sonia Adhiambo}, + year={2023}, + eprint={2305.06897}, + archivePrefix={arXiv}, + primaryClass={cs.CL} +} +``` diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/afriqa.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/afriqa.yaml new file mode 100644 index 0000000000000000000000000000000000000000..80810ca4c1195f281b6eaa9581bf420ff4582291 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/afriqa.yaml @@ -0,0 +1,13 @@ +group: afriqa +task: + - afriqa_prompt_1 + - afriqa_prompt_2 + - afriqa_prompt_3 + - afriqa_prompt_4 + - afriqa_prompt_5 +aggregate_metric_list: + - metric: acc + aggregation: mean + weight_by_size: true +metadata: + version: 1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_1/afriqa b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_1/afriqa new file mode 100644 index 0000000000000000000000000000000000000000..d9b6218e766a57309804e9514cf9d9682cf49131 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_1/afriqa @@ -0,0 +1,42 @@ +tag: + - afrobench_xqa_tasks + - afriqa_prompt_1 +dataset_kwargs: {trust_remote_code: True} +dataset_path: masakhane/afriqa-gold-passages +dataset_name: null +output_type: generate_until +test_split: test +fewshot_split: train +doc_to_target: answer_pivot +should_decontaminate: true +doc_to_decontamination_query: question_lang +generation_kwargs: + until: + - "\n" + do_sample: false + temperature: 0.0 +filter_list: + - name: remove_whitespace + filter: + - function: remove_whitespace + - function: take_first +target_delimiter: " " +metric_list: + - metric: exact_match + aggregation: mean + higher_is_better: true + ignore_case: true + ignore_punctuation: true + - "." + - "," + - "\\$" + - metric: f1 + aggregation: !function utils.f1 + higher_is_better: true + ignore_case: true + ignore_punctuation: true + - "." + - "," + - "\\$" +metadata: + version: 1.0 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_1/afriqa_bem.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_1/afriqa_bem.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a3b639a833ee4e8fb32d992538b747ab92b1f360 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_1/afriqa_bem.yaml @@ -0,0 +1,12 @@ +# Generated by utils.py +dataset_name: bem +doc_to_text: 'Your task is to answer a qestion given a context.Make sure you respond + with the shortest span containing the answer in the context. + + Question: {{question_lang}} + + Context: {{context}} + + Answer:' +include: afriqa +task: afriqa_bem_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_1/afriqa_fon.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_1/afriqa_fon.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c51196157f95e96315d0321e4b53859ef8e5ae35 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_1/afriqa_fon.yaml @@ -0,0 +1,12 @@ +# Generated by utils.py +dataset_name: fon +doc_to_text: 'Your task is to answer a qestion given a context.Make sure you respond + with the shortest span containing the answer in the context. + + Question: {{question_lang}} + + Context: {{context}} + + Answer:' +include: afriqa +task: afriqa_fon_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_1/afriqa_hau.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_1/afriqa_hau.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0536590ac3486f815ac35d5333b8a6a268fd851a --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_1/afriqa_hau.yaml @@ -0,0 +1,12 @@ +# Generated by utils.py +dataset_name: hau +doc_to_text: 'Your task is to answer a qestion given a context.Make sure you respond + with the shortest span containing the answer in the context. + + Question: {{question_lang}} + + Context: {{context}} + + Answer:' +include: afriqa +task: afriqa_hau_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_1/afriqa_ibo.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_1/afriqa_ibo.yaml new file mode 100644 index 0000000000000000000000000000000000000000..62eb71160c6cf546b274b8724cd8955c0b0e86c8 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_1/afriqa_ibo.yaml @@ -0,0 +1,12 @@ +# Generated by utils.py +dataset_name: ibo +doc_to_text: 'Your task is to answer a qestion given a context.Make sure you respond + with the shortest span containing the answer in the context. + + Question: {{question_lang}} + + Context: {{context}} + + Answer:' +include: afriqa +task: afriqa_ibo_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_1/afriqa_kin.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_1/afriqa_kin.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e632c4beef739a3e299280071229521db78cbf21 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_1/afriqa_kin.yaml @@ -0,0 +1,12 @@ +# Generated by utils.py +dataset_name: kin +doc_to_text: 'Your task is to answer a qestion given a context.Make sure you respond + with the shortest span containing the answer in the context. + + Question: {{question_lang}} + + Context: {{context}} + + Answer:' +include: afriqa +task: afriqa_kin_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_1/afriqa_swa.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_1/afriqa_swa.yaml new file mode 100644 index 0000000000000000000000000000000000000000..dbdebe14dccc3eda73ee706e3327a13a43a3aa81 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_1/afriqa_swa.yaml @@ -0,0 +1,15 @@ +# Generated by utils.py +dataset_name: swa +doc_to_text: 'Your task is to answer a qestion given a context.Make sure you respond + with the shortest span containing the answer in the context. + + Question: {{question_lang}} + + Context: {{context}} + + Answer:' +include: afriqa +fewshot_split: test +fewshot_config: + sampler: first_n +task: afriqa_swa_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_1/afriqa_twi.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_1/afriqa_twi.yaml new file mode 100644 index 0000000000000000000000000000000000000000..67ba171569e34842e2c5af86875d2495f4715421 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_1/afriqa_twi.yaml @@ -0,0 +1,12 @@ +# Generated by utils.py +dataset_name: twi +doc_to_text: 'Your task is to answer a qestion given a context.Make sure you respond + with the shortest span containing the answer in the context. + + Question: {{question_lang}} + + Context: {{context}} + + Answer:' +include: afriqa +task: afriqa_twi_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_1/afriqa_yor.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_1/afriqa_yor.yaml new file mode 100644 index 0000000000000000000000000000000000000000..51d20e43e0b4dcb6d5da0e61c22ad827085b253f --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_1/afriqa_yor.yaml @@ -0,0 +1,12 @@ +# Generated by utils.py +dataset_name: yor +doc_to_text: 'Your task is to answer a qestion given a context.Make sure you respond + with the shortest span containing the answer in the context. + + Question: {{question_lang}} + + Context: {{context}} + + Answer:' +include: afriqa +task: afriqa_yor_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_1/afriqa_zul.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_1/afriqa_zul.yaml new file mode 100644 index 0000000000000000000000000000000000000000..1c254b96e565e01750c6a838b427f926ed3f40d0 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_1/afriqa_zul.yaml @@ -0,0 +1,12 @@ +# Generated by utils.py +dataset_name: zul +doc_to_text: 'Your task is to answer a qestion given a context.Make sure you respond + with the shortest span containing the answer in the context. + + Question: {{question_lang}} + + Context: {{context}} + + Answer:' +include: afriqa +task: afriqa_zul_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_1/utils.py b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_1/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..eae1d885037da14892b39715c49e4d3aac61f06f --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_1/utils.py @@ -0,0 +1,53 @@ +import re +import string +from collections import Counter + + +def normalize_answer(s): + """ + Taken from the official evaluation script for v1.1 of the SQuAD dataset. + Lower text and remove punctuation, articles and extra whitespace. + """ + + def remove_articles(text): + return re.sub(r"\b(a|an|the)\b", " ", text) + + def white_space_fix(text): + return " ".join(text.split()) + + def remove_punc(text): + exclude = set(string.punctuation) + return "".join(ch for ch in text if ch not in exclude) + + def lower(text): + return text.lower() + + return white_space_fix(remove_articles(remove_punc(lower(s)))) + + +def f1(items): + """ + Taken from the official evaluation script for v1.1 of the SQuAD dataset. + """ + + unzipped_list = list(zip(*items)) + golds = unzipped_list[0] + preds = unzipped_list[1] + + f1_list = [] + + for i in range(len(golds)): + prediction_tokens = normalize_answer(preds[i]).split() + references_tokens = normalize_answer(golds[i]).split() + common = Counter(prediction_tokens) & Counter(references_tokens) + num_same = sum(common.values()) + if num_same == 0: + f1_score = 0 + else: + precision = 1.0 * num_same / len(prediction_tokens) + recall = 1.0 * num_same / len(references_tokens) + f1_score = (2 * precision * recall) / (precision + recall) + + f1_list.append(f1_score) + + return sum(f1_list) / len(f1_list) diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_2/afriqa b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_2/afriqa new file mode 100644 index 0000000000000000000000000000000000000000..d53ce05b168b8ffaf1325167aaee6537b9b2dbbe --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_2/afriqa @@ -0,0 +1,42 @@ +tag: + - afrobench_xqa_tasks + - afriqa_prompt_2 +dataset_kwargs: {trust_remote_code: True} +dataset_path: masakhane/afriqa-gold-passages +dataset_name: null +output_type: generate_until +test_split: test +fewshot_split: train +doc_to_target: answer_pivot +should_decontaminate: true +doc_to_decontamination_query: question_lang +generation_kwargs: + until: + - "\n" + do_sample: false + temperature: 0.0 +filter_list: + - name: remove_whitespace + filter: + - function: remove_whitespace + - function: take_first +target_delimiter: " " +metric_list: + - metric: exact_match + aggregation: mean + higher_is_better: true + ignore_case: true + ignore_punctuation: true + - "." + - "," + - "\\$" + - metric: f1 + aggregation: !function utils.f1 + higher_is_better: true + ignore_case: true + ignore_punctuation: true + - "." + - "," + - "\\$" +metadata: + version: 1.0 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_2/afriqa_bem.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_2/afriqa_bem.yaml new file mode 100644 index 0000000000000000000000000000000000000000..2469c7f434e133f0b94c41940917b17368566dfe --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_2/afriqa_bem.yaml @@ -0,0 +1,13 @@ +# Generated by utils.py +dataset_name: bem +doc_to_text: 'Your task is to answer a question given a context. The question is in + Bemba, while the context is in English or French.Make sure you respond with the + shortest span in the context that contains the answer. + + Question: {{question_lang}} + + Context: {{context}} + + Answer:' +include: afriqa +task: afriqa_bem_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_2/afriqa_fon.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_2/afriqa_fon.yaml new file mode 100644 index 0000000000000000000000000000000000000000..384db44987a074c88615164e9da85b71fad2da37 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_2/afriqa_fon.yaml @@ -0,0 +1,13 @@ +# Generated by utils.py +dataset_name: fon +doc_to_text: 'Your task is to answer a question given a context. The question is in + Fon, while the context is in English or French.Make sure you respond with the shortest + span in the context that contains the answer. + + Question: {{question_lang}} + + Context: {{context}} + + Answer:' +include: afriqa +task: afriqa_fon_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_2/afriqa_hau.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_2/afriqa_hau.yaml new file mode 100644 index 0000000000000000000000000000000000000000..40c942eced4451a620541f0cce6fe87d8f82e5cb --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_2/afriqa_hau.yaml @@ -0,0 +1,13 @@ +# Generated by utils.py +dataset_name: hau +doc_to_text: 'Your task is to answer a question given a context. The question is in + Hausa, while the context is in English or French.Make sure you respond with the + shortest span in the context that contains the answer. + + Question: {{question_lang}} + + Context: {{context}} + + Answer:' +include: afriqa +task: afriqa_hau_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_2/afriqa_ibo.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_2/afriqa_ibo.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8198795d2d1d79855d45d1c00800e56c8ad5742b --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_2/afriqa_ibo.yaml @@ -0,0 +1,13 @@ +# Generated by utils.py +dataset_name: ibo +doc_to_text: 'Your task is to answer a question given a context. The question is in + Igbo, while the context is in English or French.Make sure you respond with the shortest + span in the context that contains the answer. + + Question: {{question_lang}} + + Context: {{context}} + + Answer:' +include: afriqa +task: afriqa_ibo_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_2/afriqa_kin.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_2/afriqa_kin.yaml new file mode 100644 index 0000000000000000000000000000000000000000..7a238ae5cc6f687aeb84d88808f994b66464d8bd --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_2/afriqa_kin.yaml @@ -0,0 +1,13 @@ +# Generated by utils.py +dataset_name: kin +doc_to_text: 'Your task is to answer a question given a context. The question is in + Kinyarwanda, while the context is in English or French.Make sure you respond with + the shortest span in the context that contains the answer. + + Question: {{question_lang}} + + Context: {{context}} + + Answer:' +include: afriqa +task: afriqa_kin_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_2/afriqa_swa.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_2/afriqa_swa.yaml new file mode 100644 index 0000000000000000000000000000000000000000..4be94d07d9ed58745b6e87d38bfe927d20116137 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_2/afriqa_swa.yaml @@ -0,0 +1,16 @@ +# Generated by utils.py +dataset_name: swa +doc_to_text: 'Your task is to answer a question given a context. The question is in + Swahili, while the context is in English or French.Make sure you respond with the + shortest span in the context that contains the answer. + + Question: {{question_lang}} + + Context: {{context}} + + Answer:' +include: afriqa +fewshot_split: test +fewshot_config: + sampler: first_n +task: afriqa_swa_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_2/afriqa_twi.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_2/afriqa_twi.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f08487d0c539b519b2707cc671f7ceda3b50387d --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_2/afriqa_twi.yaml @@ -0,0 +1,13 @@ +# Generated by utils.py +dataset_name: twi +doc_to_text: 'Your task is to answer a question given a context. The question is in + Twi, while the context is in English or French.Make sure you respond with the shortest + span in the context that contains the answer. + + Question: {{question_lang}} + + Context: {{context}} + + Answer:' +include: afriqa +task: afriqa_twi_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_2/afriqa_yor.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_2/afriqa_yor.yaml new file mode 100644 index 0000000000000000000000000000000000000000..44aee11a143976b845c8cb3a9c6a1d8cf01ccf17 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_2/afriqa_yor.yaml @@ -0,0 +1,13 @@ +# Generated by utils.py +dataset_name: yor +doc_to_text: 'Your task is to answer a question given a context. The question is in + Yoruba, while the context is in English or French.Make sure you respond with the + shortest span in the context that contains the answer. + + Question: {{question_lang}} + + Context: {{context}} + + Answer:' +include: afriqa +task: afriqa_yor_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_2/afriqa_zul.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_2/afriqa_zul.yaml new file mode 100644 index 0000000000000000000000000000000000000000..99c5b18fa243c47f843862950f07e1211745f8b5 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_2/afriqa_zul.yaml @@ -0,0 +1,13 @@ +# Generated by utils.py +dataset_name: zul +doc_to_text: 'Your task is to answer a question given a context. The question is in + Zulu, while the context is in English or French.Make sure you respond with the shortest + span in the context that contains the answer. + + Question: {{question_lang}} + + Context: {{context}} + + Answer:' +include: afriqa +task: afriqa_zul_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_2/utils.py b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_2/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..eae1d885037da14892b39715c49e4d3aac61f06f --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_2/utils.py @@ -0,0 +1,53 @@ +import re +import string +from collections import Counter + + +def normalize_answer(s): + """ + Taken from the official evaluation script for v1.1 of the SQuAD dataset. + Lower text and remove punctuation, articles and extra whitespace. + """ + + def remove_articles(text): + return re.sub(r"\b(a|an|the)\b", " ", text) + + def white_space_fix(text): + return " ".join(text.split()) + + def remove_punc(text): + exclude = set(string.punctuation) + return "".join(ch for ch in text if ch not in exclude) + + def lower(text): + return text.lower() + + return white_space_fix(remove_articles(remove_punc(lower(s)))) + + +def f1(items): + """ + Taken from the official evaluation script for v1.1 of the SQuAD dataset. + """ + + unzipped_list = list(zip(*items)) + golds = unzipped_list[0] + preds = unzipped_list[1] + + f1_list = [] + + for i in range(len(golds)): + prediction_tokens = normalize_answer(preds[i]).split() + references_tokens = normalize_answer(golds[i]).split() + common = Counter(prediction_tokens) & Counter(references_tokens) + num_same = sum(common.values()) + if num_same == 0: + f1_score = 0 + else: + precision = 1.0 * num_same / len(prediction_tokens) + recall = 1.0 * num_same / len(references_tokens) + f1_score = (2 * precision * recall) / (precision + recall) + + f1_list.append(f1_score) + + return sum(f1_list) / len(f1_list) diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_3/afriqa b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_3/afriqa new file mode 100644 index 0000000000000000000000000000000000000000..79a923b1b30075d31407e804641b71339e9bedb0 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_3/afriqa @@ -0,0 +1,42 @@ +tag: + - afrobench_xqa_tasks + - afriqa_prompt_3 +dataset_kwargs: {trust_remote_code: True} +dataset_path: masakhane/afriqa-gold-passages +dataset_name: null +output_type: generate_until +test_split: test +fewshot_split: train +doc_to_target: answer_pivot +should_decontaminate: true +doc_to_decontamination_query: question_lang +generation_kwargs: + until: + - "\n" + do_sample: false + temperature: 0.0 +filter_list: + - name: remove_whitespace + filter: + - function: remove_whitespace + - function: take_first +target_delimiter: " " +metric_list: + - metric: exact_match + aggregation: mean + higher_is_better: true + ignore_case: true + ignore_punctuation: true + - "." + - "," + - "\\$" + - metric: f1 + aggregation: !function utils.f1 + higher_is_better: true + ignore_case: true + ignore_punctuation: true + - "." + - "," + - "\\$" +metadata: + version: 1.0 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_3/afriqa_bem.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_3/afriqa_bem.yaml new file mode 100644 index 0000000000000000000000000000000000000000..3af92f5a4abc656a862c701757d056b836506582 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_3/afriqa_bem.yaml @@ -0,0 +1,12 @@ +# Generated by utils.py +dataset_name: bem +doc_to_text: 'Given the context, provide the answer to the following question.Ensure + your response is concise and directly from the context. + + Question: {{question_lang}} + + Context: {{context}} + + Answer:' +include: afriqa +task: afriqa_bem_prompt_3 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_3/afriqa_fon.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_3/afriqa_fon.yaml new file mode 100644 index 0000000000000000000000000000000000000000..73c12439632863ee0cc49a84459abd0dbe4ea985 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_3/afriqa_fon.yaml @@ -0,0 +1,12 @@ +# Generated by utils.py +dataset_name: fon +doc_to_text: 'Given the context, provide the answer to the following question.Ensure + your response is concise and directly from the context. + + Question: {{question_lang}} + + Context: {{context}} + + Answer:' +include: afriqa +task: afriqa_fon_prompt_3 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_3/afriqa_hau.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_3/afriqa_hau.yaml new file mode 100644 index 0000000000000000000000000000000000000000..ff08d081971aa54fd645c9e626231e82341ecabb --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_3/afriqa_hau.yaml @@ -0,0 +1,12 @@ +# Generated by utils.py +dataset_name: hau +doc_to_text: 'Given the context, provide the answer to the following question.Ensure + your response is concise and directly from the context. + + Question: {{question_lang}} + + Context: {{context}} + + Answer:' +include: afriqa +task: afriqa_hau_prompt_3 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_3/afriqa_ibo.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_3/afriqa_ibo.yaml new file mode 100644 index 0000000000000000000000000000000000000000..12f18a0bff19f69cbfaa0fce86cf9d61ecd08136 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_3/afriqa_ibo.yaml @@ -0,0 +1,12 @@ +# Generated by utils.py +dataset_name: ibo +doc_to_text: 'Given the context, provide the answer to the following question.Ensure + your response is concise and directly from the context. + + Question: {{question_lang}} + + Context: {{context}} + + Answer:' +include: afriqa +task: afriqa_ibo_prompt_3 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_3/afriqa_kin.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_3/afriqa_kin.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e92dec41c227b0d14e2a646893b7fe4dc55e5425 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_3/afriqa_kin.yaml @@ -0,0 +1,12 @@ +# Generated by utils.py +dataset_name: kin +doc_to_text: 'Given the context, provide the answer to the following question.Ensure + your response is concise and directly from the context. + + Question: {{question_lang}} + + Context: {{context}} + + Answer:' +include: afriqa +task: afriqa_kin_prompt_3 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_3/afriqa_swa.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_3/afriqa_swa.yaml new file mode 100644 index 0000000000000000000000000000000000000000..30c574e5fa77ad65dc91d2670720945cdb67c032 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_3/afriqa_swa.yaml @@ -0,0 +1,15 @@ +# Generated by utils.py +dataset_name: swa +doc_to_text: 'Given the context, provide the answer to the following question.Ensure + your response is concise and directly from the context. + + Question: {{question_lang}} + + Context: {{context}} + + Answer:' +include: afriqa +fewshot_split: test +fewshot_config: + sampler: first_n +task: afriqa_swa_prompt_3 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_3/afriqa_twi.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_3/afriqa_twi.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b08534d9bb98b2603a405aa7d8888a8d4230a52c --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_3/afriqa_twi.yaml @@ -0,0 +1,12 @@ +# Generated by utils.py +dataset_name: twi +doc_to_text: 'Given the context, provide the answer to the following question.Ensure + your response is concise and directly from the context. + + Question: {{question_lang}} + + Context: {{context}} + + Answer:' +include: afriqa +task: afriqa_twi_prompt_3 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_3/afriqa_yor.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_3/afriqa_yor.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d3c74ce7c3a86d7df40572096bdddf75ee5321c5 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_3/afriqa_yor.yaml @@ -0,0 +1,12 @@ +# Generated by utils.py +dataset_name: yor +doc_to_text: 'Given the context, provide the answer to the following question.Ensure + your response is concise and directly from the context. + + Question: {{question_lang}} + + Context: {{context}} + + Answer:' +include: afriqa +task: afriqa_yor_prompt_3 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_3/afriqa_zul.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_3/afriqa_zul.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c54b0bd7f054dc8111c4c72e8fca55d7e7f0a4e6 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_3/afriqa_zul.yaml @@ -0,0 +1,12 @@ +# Generated by utils.py +dataset_name: zul +doc_to_text: 'Given the context, provide the answer to the following question.Ensure + your response is concise and directly from the context. + + Question: {{question_lang}} + + Context: {{context}} + + Answer:' +include: afriqa +task: afriqa_zul_prompt_3 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_3/utils.py b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_3/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..eae1d885037da14892b39715c49e4d3aac61f06f --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_3/utils.py @@ -0,0 +1,53 @@ +import re +import string +from collections import Counter + + +def normalize_answer(s): + """ + Taken from the official evaluation script for v1.1 of the SQuAD dataset. + Lower text and remove punctuation, articles and extra whitespace. + """ + + def remove_articles(text): + return re.sub(r"\b(a|an|the)\b", " ", text) + + def white_space_fix(text): + return " ".join(text.split()) + + def remove_punc(text): + exclude = set(string.punctuation) + return "".join(ch for ch in text if ch not in exclude) + + def lower(text): + return text.lower() + + return white_space_fix(remove_articles(remove_punc(lower(s)))) + + +def f1(items): + """ + Taken from the official evaluation script for v1.1 of the SQuAD dataset. + """ + + unzipped_list = list(zip(*items)) + golds = unzipped_list[0] + preds = unzipped_list[1] + + f1_list = [] + + for i in range(len(golds)): + prediction_tokens = normalize_answer(preds[i]).split() + references_tokens = normalize_answer(golds[i]).split() + common = Counter(prediction_tokens) & Counter(references_tokens) + num_same = sum(common.values()) + if num_same == 0: + f1_score = 0 + else: + precision = 1.0 * num_same / len(prediction_tokens) + recall = 1.0 * num_same / len(references_tokens) + f1_score = (2 * precision * recall) / (precision + recall) + + f1_list.append(f1_score) + + return sum(f1_list) / len(f1_list) diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_4/afriqa b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_4/afriqa new file mode 100644 index 0000000000000000000000000000000000000000..e251f1e27fab773d7fd54364ebfc870819df5d55 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_4/afriqa @@ -0,0 +1,42 @@ +tag: + - afrobench_xqa_tasks + - afriqa_prompt_4 +dataset_kwargs: {trust_remote_code: True} +dataset_path: masakhane/afriqa-gold-passages +dataset_name: null +output_type: generate_until +test_split: test +fewshot_split: train +doc_to_target: answer_pivot +should_decontaminate: true +doc_to_decontamination_query: question_lang +generation_kwargs: + until: + - "\n" + do_sample: false + temperature: 0.0 +filter_list: + - name: remove_whitespace + filter: + - function: remove_whitespace + - function: take_first +target_delimiter: " " +metric_list: + - metric: exact_match + aggregation: mean + higher_is_better: true + ignore_case: true + ignore_punctuation: true + - "." + - "," + - "\\$" + - metric: f1 + aggregation: !function utils.f1 + higher_is_better: true + ignore_case: true + ignore_punctuation: true + - "." + - "," + - "\\$" +metadata: + version: 1.0 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_4/afriqa_bem.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_4/afriqa_bem.yaml new file mode 100644 index 0000000000000000000000000000000000000000..db3d1c2a142ae6b6b1bd86678027df58c8715785 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_4/afriqa_bem.yaml @@ -0,0 +1,13 @@ +# Generated by utils.py +dataset_name: bem +doc_to_text: 'You are an AI assistant and your task is to answer the question based + on the provided context.Your answer should be the shortest span that contains the + answer within the context. + + Question: {{question_lang}} + + Context: {{context}} + + Answer:' +include: afriqa +task: afriqa_bem_prompt_4 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_4/afriqa_fon.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_4/afriqa_fon.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0c65dd07265d54a8c2aad56c563073fab7b38b49 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_4/afriqa_fon.yaml @@ -0,0 +1,13 @@ +# Generated by utils.py +dataset_name: fon +doc_to_text: 'You are an AI assistant and your task is to answer the question based + on the provided context.Your answer should be the shortest span that contains the + answer within the context. + + Question: {{question_lang}} + + Context: {{context}} + + Answer:' +include: afriqa +task: afriqa_fon_prompt_4 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_4/afriqa_hau.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_4/afriqa_hau.yaml new file mode 100644 index 0000000000000000000000000000000000000000..baeaf020b05b04e025024ac7dd4d3f6bf86caaa9 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_4/afriqa_hau.yaml @@ -0,0 +1,13 @@ +# Generated by utils.py +dataset_name: hau +doc_to_text: 'You are an AI assistant and your task is to answer the question based + on the provided context.Your answer should be the shortest span that contains the + answer within the context. + + Question: {{question_lang}} + + Context: {{context}} + + Answer:' +include: afriqa +task: afriqa_hau_prompt_4 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_4/afriqa_ibo.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_4/afriqa_ibo.yaml new file mode 100644 index 0000000000000000000000000000000000000000..6db1cc71614a9e316eeb3b12a3f0740d9cb6e671 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_4/afriqa_ibo.yaml @@ -0,0 +1,13 @@ +# Generated by utils.py +dataset_name: ibo +doc_to_text: 'You are an AI assistant and your task is to answer the question based + on the provided context.Your answer should be the shortest span that contains the + answer within the context. + + Question: {{question_lang}} + + Context: {{context}} + + Answer:' +include: afriqa +task: afriqa_ibo_prompt_4 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_4/afriqa_kin.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_4/afriqa_kin.yaml new file mode 100644 index 0000000000000000000000000000000000000000..dc8f3678207cdbc1dab076573ed7aeea02b19b92 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_4/afriqa_kin.yaml @@ -0,0 +1,13 @@ +# Generated by utils.py +dataset_name: kin +doc_to_text: 'You are an AI assistant and your task is to answer the question based + on the provided context.Your answer should be the shortest span that contains the + answer within the context. + + Question: {{question_lang}} + + Context: {{context}} + + Answer:' +include: afriqa +task: afriqa_kin_prompt_4 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_4/afriqa_swa.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_4/afriqa_swa.yaml new file mode 100644 index 0000000000000000000000000000000000000000..4fe8fbcdf2a298f9b8c58047da94c24cddc72fdc --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_4/afriqa_swa.yaml @@ -0,0 +1,16 @@ +# Generated by utils.py +dataset_name: swa +doc_to_text: 'You are an AI assistant and your task is to answer the question based + on the provided context.Your answer should be the shortest span that contains the + answer within the context. + + Question: {{question_lang}} + + Context: {{context}} + + Answer:' +include: afriqa +fewshot_split: test +fewshot_config: + sampler: first_n +task: afriqa_swa_prompt_4 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_4/afriqa_twi.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_4/afriqa_twi.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d679cd0bb6173696e7555158356b067f2eea895c --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_4/afriqa_twi.yaml @@ -0,0 +1,13 @@ +# Generated by utils.py +dataset_name: twi +doc_to_text: 'You are an AI assistant and your task is to answer the question based + on the provided context.Your answer should be the shortest span that contains the + answer within the context. + + Question: {{question_lang}} + + Context: {{context}} + + Answer:' +include: afriqa +task: afriqa_twi_prompt_4 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_4/afriqa_yor.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_4/afriqa_yor.yaml new file mode 100644 index 0000000000000000000000000000000000000000..6011dc3313ff35fe55ed204529c5d0c1503d468a --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_4/afriqa_yor.yaml @@ -0,0 +1,13 @@ +# Generated by utils.py +dataset_name: yor +doc_to_text: 'You are an AI assistant and your task is to answer the question based + on the provided context.Your answer should be the shortest span that contains the + answer within the context. + + Question: {{question_lang}} + + Context: {{context}} + + Answer:' +include: afriqa +task: afriqa_yor_prompt_4 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_4/afriqa_zul.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_4/afriqa_zul.yaml new file mode 100644 index 0000000000000000000000000000000000000000..26a6ccad3efe93ea094a64189b5efbe1dddb9734 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_4/afriqa_zul.yaml @@ -0,0 +1,13 @@ +# Generated by utils.py +dataset_name: zul +doc_to_text: 'You are an AI assistant and your task is to answer the question based + on the provided context.Your answer should be the shortest span that contains the + answer within the context. + + Question: {{question_lang}} + + Context: {{context}} + + Answer:' +include: afriqa +task: afriqa_zul_prompt_4 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_4/utils.py b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_4/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..eae1d885037da14892b39715c49e4d3aac61f06f --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_4/utils.py @@ -0,0 +1,53 @@ +import re +import string +from collections import Counter + + +def normalize_answer(s): + """ + Taken from the official evaluation script for v1.1 of the SQuAD dataset. + Lower text and remove punctuation, articles and extra whitespace. + """ + + def remove_articles(text): + return re.sub(r"\b(a|an|the)\b", " ", text) + + def white_space_fix(text): + return " ".join(text.split()) + + def remove_punc(text): + exclude = set(string.punctuation) + return "".join(ch for ch in text if ch not in exclude) + + def lower(text): + return text.lower() + + return white_space_fix(remove_articles(remove_punc(lower(s)))) + + +def f1(items): + """ + Taken from the official evaluation script for v1.1 of the SQuAD dataset. + """ + + unzipped_list = list(zip(*items)) + golds = unzipped_list[0] + preds = unzipped_list[1] + + f1_list = [] + + for i in range(len(golds)): + prediction_tokens = normalize_answer(preds[i]).split() + references_tokens = normalize_answer(golds[i]).split() + common = Counter(prediction_tokens) & Counter(references_tokens) + num_same = sum(common.values()) + if num_same == 0: + f1_score = 0 + else: + precision = 1.0 * num_same / len(prediction_tokens) + recall = 1.0 * num_same / len(references_tokens) + f1_score = (2 * precision * recall) / (precision + recall) + + f1_list.append(f1_score) + + return sum(f1_list) / len(f1_list) diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_5/afriqa b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_5/afriqa new file mode 100644 index 0000000000000000000000000000000000000000..fab00068beb951dbab88d4baa870fabfced4f820 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_5/afriqa @@ -0,0 +1,42 @@ +tag: + - afrobench_xqa_tasks + - afriqa_prompt_5 +dataset_kwargs: {trust_remote_code: True} +dataset_path: masakhane/afriqa-gold-passages +dataset_name: null +output_type: generate_until +test_split: test +fewshot_split: train +doc_to_target: answer_pivot +should_decontaminate: true +doc_to_decontamination_query: question_lang +generation_kwargs: + until: + - "\n" + do_sample: false + temperature: 0.0 +filter_list: + - name: remove_whitespace + filter: + - function: remove_whitespace + - function: take_first +target_delimiter: " " +metric_list: + - metric: exact_match + aggregation: mean + higher_is_better: true + ignore_case: true + ignore_punctuation: true + - "." + - "," + - "\\$" + - metric: f1 + aggregation: !function utils.f1 + higher_is_better: true + ignore_case: true + ignore_punctuation: true + - "." + - "," + - "\\$" +metadata: + version: 1.0 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_5/afriqa_bem.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_5/afriqa_bem.yaml new file mode 100644 index 0000000000000000000000000000000000000000..4288845d30d2e3ec1620dd1a226bf17d385322be --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_5/afriqa_bem.yaml @@ -0,0 +1,12 @@ +# Generated by utils.py +dataset_name: bem +doc_to_text: 'Using the context, find the answer to the question.Respond with the + briefest span that includes the answer from the context. + + Question: {{question_lang}} + + Context: {{context}} + + Answer:' +include: afriqa +task: afriqa_bem_prompt_5 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_5/afriqa_fon.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_5/afriqa_fon.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c234e944783b0c456b1d0a0599d43a1d569ad18f --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_5/afriqa_fon.yaml @@ -0,0 +1,12 @@ +# Generated by utils.py +dataset_name: fon +doc_to_text: 'Using the context, find the answer to the question.Respond with the + briefest span that includes the answer from the context. + + Question: {{question_lang}} + + Context: {{context}} + + Answer:' +include: afriqa +task: afriqa_fon_prompt_5 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_5/afriqa_hau.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_5/afriqa_hau.yaml new file mode 100644 index 0000000000000000000000000000000000000000..34823c9e47d2e8d614df530559c638973a02d956 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_5/afriqa_hau.yaml @@ -0,0 +1,12 @@ +# Generated by utils.py +dataset_name: hau +doc_to_text: 'Using the context, find the answer to the question.Respond with the + briefest span that includes the answer from the context. + + Question: {{question_lang}} + + Context: {{context}} + + Answer:' +include: afriqa +task: afriqa_hau_prompt_5 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_5/afriqa_ibo.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_5/afriqa_ibo.yaml new file mode 100644 index 0000000000000000000000000000000000000000..6582d2d56632e96a3d14c646a47dd7d4b55b1652 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_5/afriqa_ibo.yaml @@ -0,0 +1,12 @@ +# Generated by utils.py +dataset_name: ibo +doc_to_text: 'Using the context, find the answer to the question.Respond with the + briefest span that includes the answer from the context. + + Question: {{question_lang}} + + Context: {{context}} + + Answer:' +include: afriqa +task: afriqa_ibo_prompt_5 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_5/afriqa_kin.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_5/afriqa_kin.yaml new file mode 100644 index 0000000000000000000000000000000000000000..ed9d6517878d8ba4e29342c915648fdb48fdd45e --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_5/afriqa_kin.yaml @@ -0,0 +1,12 @@ +# Generated by utils.py +dataset_name: kin +doc_to_text: 'Using the context, find the answer to the question.Respond with the + briefest span that includes the answer from the context. + + Question: {{question_lang}} + + Context: {{context}} + + Answer:' +include: afriqa +task: afriqa_kin_prompt_5 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_5/afriqa_swa.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_5/afriqa_swa.yaml new file mode 100644 index 0000000000000000000000000000000000000000..dfcfb147f8d7789d89e0521129ba1b01f1725384 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_5/afriqa_swa.yaml @@ -0,0 +1,15 @@ +# Generated by utils.py +dataset_name: swa +doc_to_text: 'Using the context, find the answer to the question.Respond with the + briefest span that includes the answer from the context. + + Question: {{question_lang}} + + Context: {{context}} + + Answer:' +include: afriqa +fewshot_split: test +fewshot_config: + sampler: first_n +task: afriqa_swa_prompt_5 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_5/afriqa_twi.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_5/afriqa_twi.yaml new file mode 100644 index 0000000000000000000000000000000000000000..cde555cf760b159378dbc137a4759ab3c87a6b4a --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_5/afriqa_twi.yaml @@ -0,0 +1,12 @@ +# Generated by utils.py +dataset_name: twi +doc_to_text: 'Using the context, find the answer to the question.Respond with the + briefest span that includes the answer from the context. + + Question: {{question_lang}} + + Context: {{context}} + + Answer:' +include: afriqa +task: afriqa_twi_prompt_5 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_5/afriqa_yor.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_5/afriqa_yor.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c9fa17e82c271118309bf6769efff0635cf94230 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_5/afriqa_yor.yaml @@ -0,0 +1,12 @@ +# Generated by utils.py +dataset_name: yor +doc_to_text: 'Using the context, find the answer to the question.Respond with the + briefest span that includes the answer from the context. + + Question: {{question_lang}} + + Context: {{context}} + + Answer:' +include: afriqa +task: afriqa_yor_prompt_5 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_5/afriqa_zul.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_5/afriqa_zul.yaml new file mode 100644 index 0000000000000000000000000000000000000000..427e7217f4b113e28ad3efc568b35ab093eed463 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_5/afriqa_zul.yaml @@ -0,0 +1,12 @@ +# Generated by utils.py +dataset_name: zul +doc_to_text: 'Using the context, find the answer to the question.Respond with the + briefest span that includes the answer from the context. + + Question: {{question_lang}} + + Context: {{context}} + + Answer:' +include: afriqa +task: afriqa_zul_prompt_5 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_5/utils.py b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_5/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..eae1d885037da14892b39715c49e4d3aac61f06f --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/prompt_5/utils.py @@ -0,0 +1,53 @@ +import re +import string +from collections import Counter + + +def normalize_answer(s): + """ + Taken from the official evaluation script for v1.1 of the SQuAD dataset. + Lower text and remove punctuation, articles and extra whitespace. + """ + + def remove_articles(text): + return re.sub(r"\b(a|an|the)\b", " ", text) + + def white_space_fix(text): + return " ".join(text.split()) + + def remove_punc(text): + exclude = set(string.punctuation) + return "".join(ch for ch in text if ch not in exclude) + + def lower(text): + return text.lower() + + return white_space_fix(remove_articles(remove_punc(lower(s)))) + + +def f1(items): + """ + Taken from the official evaluation script for v1.1 of the SQuAD dataset. + """ + + unzipped_list = list(zip(*items)) + golds = unzipped_list[0] + preds = unzipped_list[1] + + f1_list = [] + + for i in range(len(golds)): + prediction_tokens = normalize_answer(preds[i]).split() + references_tokens = normalize_answer(golds[i]).split() + common = Counter(prediction_tokens) & Counter(references_tokens) + num_same = sum(common.values()) + if num_same == 0: + f1_score = 0 + else: + precision = 1.0 * num_same / len(prediction_tokens) + recall = 1.0 * num_same / len(references_tokens) + f1_score = (2 * precision * recall) / (precision + recall) + + f1_list.append(f1_score) + + return sum(f1_list) / len(f1_list) diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/utils.py b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..5fef58f013ff9d31da0a952e1315cc09b53c2e74 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afriqa/utils.py @@ -0,0 +1,125 @@ +import argparse +import os + +import yaml + + +class FunctionTag: + def __init__(self, value): + self.value = value + + +def prompt_func(mode, lang): + prompt_map = { + "prompt_1": "Your task is to answer a question given a context." + "Make sure you respond with the shortest span containing the answer in the context.\n" + "Question: {{question_lang}}\n" + "Context: {{context}}\n" + "Answer:", + "prompt_2": f"Your task is to answer a question given a context. The question is in {lang}, while the context is in English or French." + "Make sure you respond with the shortest span in the context that contains the answer.\n" + "Question: {{question_lang}}\n" + "Context: {{context}}\n" + "Answer:", + "prompt_3": "Given the context, provide the answer to the following question." + "Ensure your response is concise and directly from the context.\n" + "Question: {{question_lang}}\n" + "Context: {{context}}\n" + "Answer:", + "prompt_4": "You are an AI assistant and your task is to answer the question based on the provided context." + "Your answer should be the shortest span that contains the answer within the context.\n" + "Question: {{question_lang}}\n" + "Context: {{context}}\n" + "Answer:", + "prompt_5": "Using the context, find the answer to the question." + "Respond with the briefest span that includes the answer from the context.\n" + "Question: {{question_lang}}\n" + "Context: {{context}}\n" + "Answer:", + } + return prompt_map[mode] + + +def gen_lang_yamls(output_dir: str, overwrite: bool, mode: str) -> None: + """ + Generate a yaml file for each language. + + :param output_dir: The directory to output the files to. + :param overwrite: Whether to overwrite files if they already exist. + """ + err = [] + languages = { + "bem": "Bemba", + "fon": "Fon", + "hau": "Hausa", + "ibo": "Igbo", + "kin": "Kinyarwanda", + "swa": "Swahili", + "twi": "Twi", + "wol": "Wolof", + "yor": "Yoruba", + "zul": "Zulu", + } + + for lang in languages.keys(): + try: + file_name = f"afriqa_{lang}.yaml" + task_name = f"afriqa_{lang}_{mode}" + yaml_template = "afriqa" + yaml_details = { + "include": yaml_template, + "task": task_name, + "dataset_name": lang, + "doc_to_text": prompt_func(mode, languages[lang]), + } + file_path = os.path.join(output_dir, mode) + os.makedirs(file_path, exist_ok=True) + + with open( + f"{output_dir}/{mode}/{file_name}", + "w" if overwrite else "x", + encoding="utf8", + ) as f: + f.write("# Generated by utils.py\n") + yaml.dump( + yaml_details, + f, + allow_unicode=True, + ) + except FileExistsError: + err.append(file_name) + + if len(err) > 0: + raise FileExistsError( + "Files were not created because they already exist (use --overwrite flag):" + f" {', '.join(err)}" + ) + + +def main() -> None: + """Parse CLI args and generate language-specific yaml files.""" + parser = argparse.ArgumentParser() + parser.add_argument( + "--overwrite", + default=True, + action="store_true", + help="Overwrite files if they already exist", + ) + parser.add_argument( + "--output-dir", + default="./", + help="Directory to write yaml files to", + ) + parser.add_argument( + "--mode", + default="prompt_1", + choices=["prompt_1", "prompt_2", "prompt_3", "prompt_4", "prompt_5"], + help="Prompt number", + ) + args = parser.parse_args() + + gen_lang_yamls(output_dir=args.output_dir, overwrite=args.overwrite, mode=args.mode) + + +if __name__ == "__main__": + main() diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/README.md b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/README.md new file mode 100644 index 0000000000000000000000000000000000000000..99bd489e3eb2cd99eb888a0c2903a4c6259668df --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/README.md @@ -0,0 +1,58 @@ +# + +## Paper +Title: `AfriSenti: A Twitter Sentiment Analysis Benchmark for African Languages` + +Paper Link: https://aclanthology.org/2023.emnlp-main.862/ + +## Abstract +>Africa is home to over 2,000 languages from over six language families and has the highest linguistic diversity among all continents. This includes 75 languages with at least one million speakers each. Yet, there is little NLP research conducted on African languages. Crucial in enabling such research is the availability of high-quality annotated datasets. In this paper, we introduce AfriSenti, a sentiment analysis benchmark that contains a total of >110,000 tweets in 14 African languages (Amharic, Algerian Arabic, Hausa, Igbo, Kinyarwanda, Moroccan Arabic, Mozambican Portuguese, Nigerian Pidgin, Oromo, Swahili, Tigrinya, Twi, Xitsonga, and Yoruba) from four language families. The tweets were annotated by native speakers and used in the AfriSenti-SemEval shared task (with over 200 participants, see website: https://afrisenti-semeval.github.io). We describe the data collection methodology, annotation process, and the challenges we dealt with when curating each dataset. We further report baseline experiments conducted on the AfriSenti datasets and discuss their usefulness. + +HomePage: https://github.com/afrisenti-semeval/afrisent-semeval-2023 + +### Citation + +``` +@inproceedings{muhammad-etal-2023-afrisenti, + title = "{A}fri{S}enti: A {T}witter Sentiment Analysis Benchmark for {A}frican Languages", + author = "Muhammad, Shamsuddeen Hassan and + Abdulmumin, Idris and + Ayele, Abinew Ali and + Ousidhoum, Nedjma and + Adelani, David Ifeoluwa and + Yimam, Seid Muhie and + Ahmad, Ibrahim Sa'id and + Beloucif, Meriem and + Mohammad, Saif M. and + Ruder, Sebastian and + Hourrane, Oumaima and + Brazdil, Pavel and + Jorge, Alipio and + Ali, Felermino D{\'a}rio M{\'a}rio Ant{\'o}nio and + David, Davis and + Osei, Salomey and + Shehu Bello, Bello and + Ibrahim, Falalu and + Gwadabe, Tajuddeen and + Rutunda, Samuel and + Belay, Tadesse and + Messelle, Wendimu Baye and + Balcha, Hailu Beshada and + Chala, Sisay Adugna and + Gebremichael, Hagos Tesfahun and + Opoku, Bernard and + Arthur, Stephen", + editor = "Bouamor, Houda and + Pino, Juan and + Bali, Kalika", + booktitle = "Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing", + month = dec, + year = "2023", + address = "Singapore", + publisher = "Association for Computational Linguistics", + url = "https://aclanthology.org/2023.emnlp-main.862/", + doi = "10.18653/v1/2023.emnlp-main.862", + pages = "13968--13981", + abstract = "Africa is home to over 2,000 languages from over six language families and has the highest linguistic diversity among all continents. This includes 75 languages with at least one million speakers each. Yet, there is little NLP research conducted on African languages. Crucial in enabling such research is the availability of high-quality annotated datasets. In this paper, we introduce AfriSenti, a sentiment analysis benchmark that contains a total of {\ensuremath{>}}110,000 tweets in 14 African languages (Amharic, Algerian Arabic, Hausa, Igbo, Kinyarwanda, Moroccan Arabic, Mozambican Portuguese, Nigerian Pidgin, Oromo, Swahili, Tigrinya, Twi, Xitsonga, and Yoruba) from four language families. The tweets were annotated by native speakers and used in the AfriSenti-SemEval shared task (with over 200 participants, see website: https://afrisenti-semeval.github.io). We describe the data collection methodology, annotation process, and the challenges we dealt with when curating each dataset. We further report baseline experiments conducted on the AfriSenti datasets and discuss their usefulness." +} +``` diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/afrisenti.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/afrisenti.yaml new file mode 100644 index 0000000000000000000000000000000000000000..36a1efdb3033e70060251e346847b73fd9de2f60 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/afrisenti.yaml @@ -0,0 +1,13 @@ +group: afrisenti +task: + - afrisenti_prompt_1 + - afrisenti_prompt_2 + - afrisenti_prompt_3 + - afrisenti_prompt_4 + - afrisenti_prompt_5 +aggregate_metric_list: + - metric: acc + aggregation: mean + weight_by_size: true +metadata: + version: 1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/fewshot.sh b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/fewshot.sh new file mode 100644 index 0000000000000000000000000000000000000000..428d455b65ac917efee1810a68626f36e777e2d9 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/fewshot.sh @@ -0,0 +1,109 @@ +lm_eval --model hf \ + --model_args pretrained=masakhane/African-ultrachat-alpaca \ + --tasks afrimmlu_direct_amh,afrimmlu_direct_eng,afrimmlu_direct_ewe,afrimmlu_direct_fra,afrimmlu_direct_hau,afrimmlu_direct_ibo,afrimmlu_direct_kin,afrimmlu_direct_lin,afrimmlu_direct_lug,afrimmlu_direct_orm,afrimmlu_direct_sna,afrimmlu_direct_sot,afrimmlu_direct_twi,afrimmlu_direct_wol,afrimmlu_direct_xho,afrimmlu_direct_yor,afrimmlu_direct_zul \ + --device cuda:0 \ + --batch_size 1 \ + --num_fewshot 0 \ + --verbosity DEBUG \ + --wandb_args project=afrimmlu + + +lm_eval --model hf \ + --model_args pretrained=bigscience/mt0-small,parallelize=true \ + --tasks afrisenti_amh_prompt_1,afrisenti_arq_prompt_1,afrisenti_ary_prompt_1,afrisenti_hau_prompt_1,afrisenti_ibo_prompt_1,afrisenti_kin_prompt_1,afrisenti_orm_prompt_1,afrisenti_pcm_prompt_1,afrisenti_por_prompt_1,afrisenti_swa_prompt_1,afrisenti_tir_prompt_1,afrisenti_tso_prompt_1,afrisenti_twi_prompt_1,afrisenti_yor_prompt_1\ + --device cuda:0 \ + --batch_size 1 \ + --num_fewshot 0 \ + --verbosity DEBUG \ + --limit 5 + + +lm_eval --model hf \ + --model_args pretrained=bigscience/mt0-xxl,parallelize=true \ + --tasks afrisenti_amh_prompt_1,afrisenti_arq_prompt_1,afrisenti_ary_prompt_1,afrisenti_hau_prompt_1,afrisenti_ibo_prompt_1,afrisenti_kin_prompt_1,afrisenti_orm_prompt_1,afrisenti_pcm_prompt_1,afrisenti_por_prompt_1,afrisenti_swa_prompt_1,afrisenti_tir_prompt_1,afrisenti_tso_prompt_1,afrisenti_twi_prompt_1,afrisenti_yor_prompt_1\ + --batch_size 128 \ + --num_fewshot 0 \ + --verbosity DEBUG + +lm_eval --model hf \ + --model_args pretrained=google/gemma-2-27b-it,parallelize=true,trust_remote_code=True \ + --tasks afriqa_wol_prompt_2\ + --batch_size 1 \ + --device 'cuda' \ + --num_fewshot 5 \ + --verbosity DEBUG \ + --output_path './afriqa_results/' \ + --log_samples + +lm_eval --model vllm \ + --model_args pretrained=meta-llama/Llama-2-7b-chat-hf,tensor_parallel_size=2,dtype='auto',gpu_memory_utilization=0.8,data_parallel_size=1 \ + --tasks masakhapos_pcm_prompt_1,masakhapos_pcm_prompt_2,masakhapos_pcm_prompt_3,masakhapos_pcm_prompt_4,masakhapos_pcm_prompt_5 \ + --batch_size 'auto' \ + --device 'cuda' \ + --num_fewshot 0 \ + --verbosity DEBUG \ + --limit 2 + + +lm_eval --model vllm \ + --model_args pretrained=meta-llama/Llama-2-7b-chat-hf,tensor_parallel_size=2,dtype='auto',gpu_memory_utilization=0.8,data_parallel_size=1 \ + --tasks masakhapos_pcm_prompt_1,masakhapos_pcm_prompt_2,masakhapos_pcm_prompt_3,masakhapos_bam_prompt_2,masakhapos_bbj_prompt_3 \ + --batch_size 'auto' \ + --device 'cuda' \ + --num_fewshot 0 \ + --verbosity DEBUG + +lm_eval --model vllm \ + --model_args pretrained=google/gemma-1.1-7b-it,tensor_parallel_size=2,dtype='auto',gpu_memory_utilization=0.8,data_parallel_size=1 \ + --tasks masakhaner_pcm_prompt_1\ + --batch_size 'auto' \ + --device 'cuda' \ + --num_fewshot 0 \ + --verbosity DEBUG \ + --limit 5 + +lm_eval --model vllm \ + --model_args pretrained=google/gemma-2-9b-it,tensor_parallel_size=2,dtype='auto',gpu_memory_utilization=0.8,data_parallel_size=1 \ + --tasks masakhaner_pcm_prompt_1,masakhaner_pcm_prompt_2,masakhaner_pcm_prompt_3,masakhaner_pcm_prompt_4,masakhaner_pcm_prompt_5\ + --batch_size 'auto' \ + --device 'cuda' \ + --num_fewshot 0 \ + --verbosity DEBUG \ + --limit 5 + +lm_eval --model vllm \ + --model_args pretrained=google/gemma-1.1-7b-it,tensor_parallel_size=2,dtype='auto',gpu_memory_utilization=0.8,data_parallel_size=1 \ + --tasks flores_eng_Latn-fuv_Latn_prompt_1,flores_eng_Latn-fuv_Latn_prompt_2,flores_eng_Latn-fuv_Latn_prompt_3,flores_fuv_Latn-eng_Latn_prompt_1,flores_fuv_Latn-eng_Latn_prompt_2,flores_fuv_Latn-eng_Latn_prompt_3 \ + --batch_size 'auto' \ + --device 'cuda' \ + --num_fewshot 0 \ + --verbosity DEBUG \ + --limit 2 + +lm_eval --model vllm \ + --model_args pretrained=google/gemma-2-27b-it,tensor_parallel_size=2,dtype='auto',gpu_memory_utilization=0.9,data_parallel_size=1 \ + --tasks masakhapos_twi_prompt_3,masakhapos_wol_prompt_3,masakhapos_xho_prompt_3,masakhapos_yor_prompt_3,masakhapos_zul_prompt_3\ + --batch_size 'auto' \ + --num_fewshot 5 \ + --verbosity DEBUG \ + --output_path './masakhapos_results/' \ + --log_samples + +lm_eval --model hf \ + --model_args pretrained=bigscience/mt0-small,parallelize=true \ + --tasks injongointent_amh_prompt_1,injongointent_eng_prompt_1,injongointent_yor_prompt_1,injongointent_ibo_prompt_1,injongointent_wol_prompt_1\ + --device 'mps' \ + --batch_size 1 \ + --num_fewshot 0 \ + --verbosity DEBUG \ + --limit 5 + +lm_eval --model hf \ + --model_args pretrained=google/gemma-3-27b-it,parallelize=true \ + --tasks afrobench_sentiment_tasks\ + --device 'cuda' \ + --batch_size 1 \ + --num_fewshot 0 \ + --verbosity DEBUG \ + --output_path './senti_results/' \ + --log_samples diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/afrisenti b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/afrisenti new file mode 100644 index 0000000000000000000000000000000000000000..69ef6b2bc08bbc198e2c6610c7c40041db4d20a4 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/afrisenti @@ -0,0 +1,41 @@ +tag: + - afrobench_sentiment_tasks + - afrisenti_prompt_1 +task: null +dataset_path: masakhane/afrisenti +dataset_name: null +dataset_kwargs: {trust_remote_code: True} +output_type: multiple_choice +validation_split: validation +test_split: test +fewshot_split: train +doc_to_text: 'Does this statement; "{{tweet}}" have a Neutral, Positive or Negative sentiment? Labels only' +doc_to_target: label +doc_to_choice: + - "negative" + - "positive" + - "neutral" +should_decontaminate: true +doc_to_decontamination_query: tweet +metric_list: + - metric: f1 + aggregation: !function utils.weighted_f1_score + # aggregation: mean + average: weighted + hf_evaluate: true + higher_is_better: True + ignore_case: true + ignore_punctuation: true + regexes_to_ignore: + - "," + - "\\$" + - metric: acc + aggregation: mean + higher_is_better: true + ignore_case: true + ignore_punctuation: true + regexes_to_ignore: + - "," + - "\\$" +metadata: + version: 1.0 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/afrisenti_amh.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/afrisenti_amh.yaml new file mode 100644 index 0000000000000000000000000000000000000000..7eefbe867360070e0701a558c124ad4ad7da786a --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/afrisenti_amh.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: amh +include: afrisenti +task: afrisenti_amh_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/afrisenti_arq.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/afrisenti_arq.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8b2e2522d94e2d0da8cb9efec72d225c7161f8e7 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/afrisenti_arq.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: arq +include: afrisenti +task: afrisenti_arq_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/afrisenti_ary.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/afrisenti_ary.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8f9ef3f20d654a0e97e40a9dd4e3d6bd2e7d949b --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/afrisenti_ary.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: ary +include: afrisenti +task: afrisenti_ary_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/afrisenti_hau.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/afrisenti_hau.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f0ab9071abbc0211b8048743db43347bb5df1583 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/afrisenti_hau.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: hau +include: afrisenti +task: afrisenti_hau_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/afrisenti_ibo.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/afrisenti_ibo.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b0176d08764dae9a5fd8af57dc903b6b55ab0124 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/afrisenti_ibo.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: ibo +include: afrisenti +task: afrisenti_ibo_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/afrisenti_kin.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/afrisenti_kin.yaml new file mode 100644 index 0000000000000000000000000000000000000000..75bb717a6e22d931404d2c7cfc919cfa8f99453d --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/afrisenti_kin.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: kin +include: afrisenti +task: afrisenti_kin_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/afrisenti_orm.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/afrisenti_orm.yaml new file mode 100644 index 0000000000000000000000000000000000000000..65c63b06fbb721d59033e5f748c6899270e92831 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/afrisenti_orm.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: orm +include: afrisenti +task: afrisenti_orm_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/afrisenti_pcm.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/afrisenti_pcm.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0f24fe9fc01cc162294aa1b387676591450b2d39 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/afrisenti_pcm.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: pcm +include: afrisenti +task: afrisenti_pcm_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/afrisenti_por.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/afrisenti_por.yaml new file mode 100644 index 0000000000000000000000000000000000000000..1e1b4cd60a1533ab6804624c8e967b46c37a69be --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/afrisenti_por.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: por +include: afrisenti +task: afrisenti_por_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/afrisenti_swa.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/afrisenti_swa.yaml new file mode 100644 index 0000000000000000000000000000000000000000..3386948ccf5eef56d293cc39d0811ab06e1a5127 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/afrisenti_swa.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: swa +include: afrisenti +task: afrisenti_swa_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/afrisenti_tir.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/afrisenti_tir.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c4942628e8115f58a047d7819052b27cc50883e9 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/afrisenti_tir.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: tir +include: afrisenti +task: afrisenti_tir_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/afrisenti_tso.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/afrisenti_tso.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d500693270946b6581020bc68a38612bdfd4f033 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/afrisenti_tso.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: tso +include: afrisenti +task: afrisenti_tso_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/afrisenti_twi.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/afrisenti_twi.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a68bb23dcaeedcc6c97768c42a1e49c26b425e40 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/afrisenti_twi.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: twi +include: afrisenti +task: afrisenti_twi_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/afrisenti_yor.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/afrisenti_yor.yaml new file mode 100644 index 0000000000000000000000000000000000000000..fda98c2c82c6e323eae8c96843e657e07a4d9665 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/afrisenti_yor.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: yor +include: afrisenti +task: afrisenti_yor_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/run.sh b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/run.sh new file mode 100644 index 0000000000000000000000000000000000000000..50d1a1338f87330219dd4c6f79fa85ef918bb21c --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/run.sh @@ -0,0 +1,34 @@ +#!/bin/bash + +models=( + + "google/gemma-1.1-7b-it" + "CohereForAI/aya-101" + "meta-llama/Llama-2-7b-chat-hf" + "meta-llama/Meta-Llama-3-8B-Instruct" + "google/gemma-2-9b-it" + "bigscience/mt0-xxl" + "google/gemma-2-27b-it" + "meta-llama/Meta-Llama-3-70B-Instruct" +) +task=afrisenti_amh_prompt_1,afrisenti_arq_prompt_1,afrisenti_ary_prompt_1,afrisenti_hau_prompt_1,afrisenti_ibo_prompt_1,afrisenti_kin_prompt_1,afrisenti_pcm_prompt_1,afrisenti_por_prompt_1,afrisenti_swa_prompt_1,afrisenti_tir_prompt_1,afrisenti_tso_prompt_1,afrisenti_twi_prompt_1,afrisenti_yor_prompt_1 + +for model in "${models[@]}" +do + echo "Evaluating model: $model" + for fewshot in 0 5 + do + export OUTPUT_DIR=results/$fewshot + + mkdir -p "$OUTPUT_DIR" + + lm_eval --model hf \ + --model_args "pretrained=${model}" \ + --tasks $task\ + --device cuda:0 \ + --batch_size 16 \ + --output_path "$OUTPUT_DIR" \ + --num_fewshot $fewshot \ + --verbosity DEBUG + done +done diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/utils.py b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..3e735e2deb1f9c53152c072615aebe8ba3acb90b --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/utils.py @@ -0,0 +1 @@ +from lm_eval.utils import weighted_f1_score diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/xx.py b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/xx.py new file mode 100644 index 0000000000000000000000000000000000000000..ca0e325e526e33dfd19ba03a93652244977fc119 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_1/xx.py @@ -0,0 +1,13 @@ +from datasets import load_dataset + + +# ['amh', 'hau', 'ibo', 'arq', 'ary', 'yor', 'por', 'twi', 'tso', 'tir', 'orm', 'pcm', 'kin', 'swa'] + +data = load_dataset("masakhane/afrisenti", "pcm", trust_remote_code=True) +print(data) +print(data["test"][:5]) +# +# ['Naija', 'Pipo', 'wey', 'dey', 'for', 'inside', 'social', 'Media', 'sef', 'don', 'put', 'hand', 'for', 'ear', 'give', +# 'federal', 'goment', 'and', 'polical', 'leader', 'dem', 'ova', 'di', 'kilin', '.'] +# +# [6, 0, 14, 17, 2, 2, 6, 0, 7, 17, 16, 0, 2, 0, 16, 0, 0, 9, 0, 0, 11, 2, 8, 0, 1] diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/afrisenti b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/afrisenti new file mode 100644 index 0000000000000000000000000000000000000000..879f2826c3f26025fcb5e41342f86ef3f9c6c677 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/afrisenti @@ -0,0 +1,39 @@ +tag: + - afrobench_sentiment_tasks + - afrisent_prompt_2 +dataset_path: masakhane/afrisenti +dataset_name: null +dataset_kwargs: {trust_remote_code: True} +output_type: multiple_choice +validation_split: validation +test_split: test +fewshot_split: train +doc_to_target: label +doc_to_choice: + - "negative" + - "positive" + - "neutral" +should_decontaminate: true +doc_to_decontamination_query: 'text: {{tweet}} \nlabel: ' +metric_list: + - metric: f1 + aggregation: !function utils.weighted_f1_score + # aggregation: mean + average: weighted + hf_evaluate: true + higher_is_better: True + ignore_case: true + ignore_punctuation: true + regexes_to_ignore: + - "," + - "\\$" + - metric: acc + aggregation: mean + higher_is_better: true + ignore_case: true + ignore_punctuation: true + regexes_to_ignore: + - "," + - "\\$" +metadata: + version: 1.0 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/afrisenti_amh.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/afrisenti_amh.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d97b2c25787d9338546dba3707afcda31ad31269 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/afrisenti_amh.yaml @@ -0,0 +1,6 @@ +# Generated by utils.py +dataset_name: amh +doc_to_text: Does this Amharic statement; '{{tweet}}' have a Neutral, Positive or + Negative sentiment? Labels only +include: afrisenti +task: afrisenti_amh_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/afrisenti_arq.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/afrisenti_arq.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c61e310dee3a9f97557aaaa8d465ce329ba29610 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/afrisenti_arq.yaml @@ -0,0 +1,6 @@ +# Generated by utils.py +dataset_name: arq +doc_to_text: Does this Algerian Arabic statement; '{{tweet}}' have a Neutral, Positive + or Negative sentiment? Labels only +include: afrisenti +task: afrisenti_arq_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/afrisenti_ary.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/afrisenti_ary.yaml new file mode 100644 index 0000000000000000000000000000000000000000..4e76d385b3dc5ee5bbbe817336ec9d78feef7eb7 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/afrisenti_ary.yaml @@ -0,0 +1,6 @@ +# Generated by utils.py +dataset_name: ary +doc_to_text: Does this Moroccan Arabic statement; '{{tweet}}' have a Neutral, Positive + or Negative sentiment? Labels only +include: afrisenti +task: afrisenti_ary_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/afrisenti_hau.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/afrisenti_hau.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f7b0ccb2811b30acc98fe594af33b0a38fb2a88b --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/afrisenti_hau.yaml @@ -0,0 +1,6 @@ +# Generated by utils.py +dataset_name: hau +doc_to_text: Does this Hausa statement; '{{tweet}}' have a Neutral, Positive or Negative + sentiment? Labels only +include: afrisenti +task: afrisenti_hau_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/afrisenti_ibo.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/afrisenti_ibo.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d4d6c6c8094bb5b41e5488c972f2702705c9f3d5 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/afrisenti_ibo.yaml @@ -0,0 +1,6 @@ +# Generated by utils.py +dataset_name: ibo +doc_to_text: Does this Igbo statement; '{{tweet}}' have a Neutral, Positive or Negative + sentiment? Labels only +include: afrisenti +task: afrisenti_ibo_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/afrisenti_kin.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/afrisenti_kin.yaml new file mode 100644 index 0000000000000000000000000000000000000000..5067b9fb75cd1579a8bc915a3a010696ca60b177 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/afrisenti_kin.yaml @@ -0,0 +1,6 @@ +# Generated by utils.py +dataset_name: kin +doc_to_text: Does this Kinyarwanda statement; '{{tweet}}' have a Neutral, Positive + or Negative sentiment? Labels only +include: afrisenti +task: afrisenti_kin_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/afrisenti_orm.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/afrisenti_orm.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e8abbbfbd73687a4249238a3f2ad988b85634531 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/afrisenti_orm.yaml @@ -0,0 +1,6 @@ +# Generated by utils.py +dataset_name: orm +doc_to_text: Does this Oromo statement; '{{tweet}}' have a Neutral, Positive or Negative + sentiment? Labels only +include: afrisenti +task: afrisenti_orm_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/afrisenti_pcm.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/afrisenti_pcm.yaml new file mode 100644 index 0000000000000000000000000000000000000000..4dd98925299e9b872d49bdb9bea0ebd2b48d1ec7 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/afrisenti_pcm.yaml @@ -0,0 +1,6 @@ +# Generated by utils.py +dataset_name: pcm +doc_to_text: Does this Nigerian Pidgin statement; '{{tweet}}' have a Neutral, Positive + or Negative sentiment? Labels only +include: afrisenti +task: afrisenti_pcm_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/afrisenti_por.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/afrisenti_por.yaml new file mode 100644 index 0000000000000000000000000000000000000000..4b8beecff946f4764becafb103f890ea924926a9 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/afrisenti_por.yaml @@ -0,0 +1,6 @@ +# Generated by utils.py +dataset_name: por +doc_to_text: Does this Mozambique Portuguese statement; '{{tweet}}' have a Neutral, + Positive or Negative sentiment? Labels only +include: afrisenti +task: afrisenti_por_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/afrisenti_swa.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/afrisenti_swa.yaml new file mode 100644 index 0000000000000000000000000000000000000000..496da1a1d1e2b4fcfa004e918af85e7321a1ed29 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/afrisenti_swa.yaml @@ -0,0 +1,6 @@ +# Generated by utils.py +dataset_name: swa +doc_to_text: Does this Swahili statement; '{{tweet}}' have a Neutral, Positive or + Negative sentiment? Labels only +include: afrisenti +task: afrisenti_swa_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/afrisenti_tir.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/afrisenti_tir.yaml new file mode 100644 index 0000000000000000000000000000000000000000..3899c992ed3180d76f2ff677c148026da6d5e9a1 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/afrisenti_tir.yaml @@ -0,0 +1,6 @@ +# Generated by utils.py +dataset_name: tir +doc_to_text: Does this Tigrinya statement; '{{tweet}}' have a Neutral, Positive or + Negative sentiment? Labels only +include: afrisenti +task: afrisenti_tir_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/afrisenti_tso.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/afrisenti_tso.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b371b7479489bd5387d2d18cd7f2edab8496dc00 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/afrisenti_tso.yaml @@ -0,0 +1,6 @@ +# Generated by utils.py +dataset_name: tso +doc_to_text: Does this Xithonga statement; '{{tweet}}' have a Neutral, Positive or + Negative sentiment? Labels only +include: afrisenti +task: afrisenti_tso_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/afrisenti_twi.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/afrisenti_twi.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c985efc4d32d30ae7f18ed5aac13c97d4dbe112b --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/afrisenti_twi.yaml @@ -0,0 +1,6 @@ +# Generated by utils.py +dataset_name: twi +doc_to_text: Does this Twi statement; '{{tweet}}' have a Neutral, Positive or Negative + sentiment? Labels only +include: afrisenti +task: afrisenti_twi_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/afrisenti_yor.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/afrisenti_yor.yaml new file mode 100644 index 0000000000000000000000000000000000000000..78932ed4cfe5f88bc91a6a0d26eb8c33a71c1ecb --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/afrisenti_yor.yaml @@ -0,0 +1,6 @@ +# Generated by utils.py +dataset_name: yor +doc_to_text: Does this Yoruba statement; '{{tweet}}' have a Neutral, Positive or Negative + sentiment? Labels only +include: afrisenti +task: afrisenti_yor_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/run.sh b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/run.sh new file mode 100644 index 0000000000000000000000000000000000000000..48797912512124c9c5287dcdd654e5fa04a029b0 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/run.sh @@ -0,0 +1,33 @@ +#!/bin/bash + +models=( + + "google/gemma-1.1-7b-it" + "CohereForAI/aya-101" + "meta-llama/Llama-2-7b-chat-hf" + "meta-llama/Meta-Llama-3-8B-Instruct" + "google/gemma-2-9b-it" + "bigscience/mt0-xxl" + "google/gemma-2-27b-it" + "meta-llama/Meta-Llama-3-70B-Instruct" +) + +for model in "${models[@]}" +do + echo "Evaluating model: $model" + for fewshot in 0 5 + do + export OUTPUT_DIR=./results/$fewshot + + mkdir -p "$OUTPUT_DIR" + + lm_eval --model hf \ + --model_args "pretrained=${model},parallelize: true" \ + --tasks afribench\ + --batch_size 256 \ + --output_path "$OUTPUT_DIR" \ + --num_fewshot $fewshot \ + --verbosity DEBUG \ + --limit 2 + done +done diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/utils.py b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..3e735e2deb1f9c53152c072615aebe8ba3acb90b --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_2/utils.py @@ -0,0 +1 @@ +from lm_eval.utils import weighted_f1_score diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_3/afrisenti b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_3/afrisenti new file mode 100644 index 0000000000000000000000000000000000000000..53cb77771f2cc6622fa4c67ea5ea20485df761d6 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_3/afrisenti @@ -0,0 +1,39 @@ +tag: + - afrobench_sentiment_tasks + - afrisenti_prompt_3 +dataset_path: masakhane/afrisenti +dataset_name: null +dataset_kwargs: {trust_remote_code: True} +output_type: multiple_choice +validation_split: validation +test_split: test +fewshot_split: train +doc_to_target: label +doc_to_choice: + - "negative" + - "positive" + - "neutral" +should_decontaminate: true +doc_to_decontamination_query: 'text: {{tweet}} \nlabel: ' +metric_list: + - metric: f1 + aggregation: !function utils.weighted_f1_score + # aggregation: mean + average: weighted + hf_evaluate: true + higher_is_better: True + ignore_case: true + ignore_punctuation: true + regexes_to_ignore: + - "," + - "\\$" + - metric: acc + aggregation: mean + higher_is_better: true + ignore_case: true + ignore_punctuation: true + regexes_to_ignore: + - "," + - "\\$" +metadata: + version: 1.0 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_3/afrisenti_arq.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_3/afrisenti_arq.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0b90f690e93f249e8c4668bb74a667ff19a39247 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_3/afrisenti_arq.yaml @@ -0,0 +1,8 @@ +# Generated by utils.py +dataset_name: arq +doc_to_text: "You are an assistant able to detect sentiments in tweets. \n\nGiven\ + \ the sentiment labels Neutral, Positive or Negative; what is the sentiment of the\ + \ Algerian Arabic statement below? Return only the labels. \n\ntext: {{tweet}} \n\ + label:" +include: afrisenti +task: afrisenti_arq_prompt_3 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_3/afrisenti_ary.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_3/afrisenti_ary.yaml new file mode 100644 index 0000000000000000000000000000000000000000..ba11ee3e5146db8bfef7680becfb95a6f0a9b6aa --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_3/afrisenti_ary.yaml @@ -0,0 +1,8 @@ +# Generated by utils.py +dataset_name: ary +doc_to_text: "You are an assistant able to detect sentiments in tweets. \n\nGiven\ + \ the sentiment labels Neutral, Positive or Negative; what is the sentiment of the\ + \ Moroccan Arabic statement below? Return only the labels. \n\ntext: {{tweet}} \n\ + label:" +include: afrisenti +task: afrisenti_ary_prompt_3 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_3/afrisenti_hau.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_3/afrisenti_hau.yaml new file mode 100644 index 0000000000000000000000000000000000000000..4f4e6b3fb3252929fcd2d0f30240ca8d4553a009 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_3/afrisenti_hau.yaml @@ -0,0 +1,7 @@ +# Generated by utils.py +dataset_name: hau +doc_to_text: "You are an assistant able to detect sentiments in tweets. \n\nGiven\ + \ the sentiment labels Neutral, Positive or Negative; what is the sentiment of the\ + \ Hausa statement below? Return only the labels. \n\ntext: {{tweet}} \nlabel:" +include: afrisenti +task: afrisenti_hau_prompt_3 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_3/afrisenti_kin.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_3/afrisenti_kin.yaml new file mode 100644 index 0000000000000000000000000000000000000000..52d84b2684f9a071c0dc6b5889ad790e3116b0fc --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_3/afrisenti_kin.yaml @@ -0,0 +1,7 @@ +# Generated by utils.py +dataset_name: kin +doc_to_text: "You are an assistant able to detect sentiments in tweets. \n\nGiven\ + \ the sentiment labels Neutral, Positive or Negative; what is the sentiment of the\ + \ Kinyarwanda statement below? Return only the labels. \n\ntext: {{tweet}} \nlabel:" +include: afrisenti +task: afrisenti_kin_prompt_3 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_3/afrisenti_orm.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_3/afrisenti_orm.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e2d524bfd781b37cb156ce09fd7fc5aae493392b --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_3/afrisenti_orm.yaml @@ -0,0 +1,7 @@ +# Generated by utils.py +dataset_name: orm +doc_to_text: "You are an assistant able to detect sentiments in tweets. \n\nGiven\ + \ the sentiment labels Neutral, Positive or Negative; what is the sentiment of the\ + \ Oromo statement below? Return only the labels. \n\ntext: {{tweet}} \nlabel:" +include: afrisenti +task: afrisenti_orm_prompt_3 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_3/afrisenti_pcm.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_3/afrisenti_pcm.yaml new file mode 100644 index 0000000000000000000000000000000000000000..eb0ac8ff3bc1b3e0c8efa9fe9afdc40ea0c0690a --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_3/afrisenti_pcm.yaml @@ -0,0 +1,8 @@ +# Generated by utils.py +dataset_name: pcm +doc_to_text: "You are an assistant able to detect sentiments in tweets. \n\nGiven\ + \ the sentiment labels Neutral, Positive or Negative; what is the sentiment of the\ + \ Nigerian Pidgin statement below? Return only the labels. \n\ntext: {{tweet}} \n\ + label:" +include: afrisenti +task: afrisenti_pcm_prompt_3 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_3/afrisenti_por.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_3/afrisenti_por.yaml new file mode 100644 index 0000000000000000000000000000000000000000..821a4355b044844d3608c01950b949ce5f292ba2 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_3/afrisenti_por.yaml @@ -0,0 +1,8 @@ +# Generated by utils.py +dataset_name: por +doc_to_text: "You are an assistant able to detect sentiments in tweets. \n\nGiven\ + \ the sentiment labels Neutral, Positive or Negative; what is the sentiment of the\ + \ Mozambique Portuguese statement below? Return only the labels. \n\ntext: {{tweet}}\ + \ \nlabel:" +include: afrisenti +task: afrisenti_por_prompt_3 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_3/afrisenti_swa.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_3/afrisenti_swa.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d8e92842e01b61cc815fb48f7a390c6f13587e18 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_3/afrisenti_swa.yaml @@ -0,0 +1,7 @@ +# Generated by utils.py +dataset_name: swa +doc_to_text: "You are an assistant able to detect sentiments in tweets. \n\nGiven\ + \ the sentiment labels Neutral, Positive or Negative; what is the sentiment of the\ + \ Swahili statement below? Return only the labels. \n\ntext: {{tweet}} \nlabel:" +include: afrisenti +task: afrisenti_swa_prompt_3 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_3/afrisenti_tso.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_3/afrisenti_tso.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8355035e963a13e2da541c65d4f778c5f0d46a58 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_3/afrisenti_tso.yaml @@ -0,0 +1,7 @@ +# Generated by utils.py +dataset_name: tso +doc_to_text: "You are an assistant able to detect sentiments in tweets. \n\nGiven\ + \ the sentiment labels Neutral, Positive or Negative; what is the sentiment of the\ + \ Xithonga statement below? Return only the labels. \n\ntext: {{tweet}} \nlabel:" +include: afrisenti +task: afrisenti_tso_prompt_3 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_3/afrisenti_twi.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_3/afrisenti_twi.yaml new file mode 100644 index 0000000000000000000000000000000000000000..98809176e9693cd65ec711fb81739fdbe0030e70 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_3/afrisenti_twi.yaml @@ -0,0 +1,7 @@ +# Generated by utils.py +dataset_name: twi +doc_to_text: "You are an assistant able to detect sentiments in tweets. \n\nGiven\ + \ the sentiment labels Neutral, Positive or Negative; what is the sentiment of the\ + \ Twi statement below? Return only the labels. \n\ntext: {{tweet}} \nlabel:" +include: afrisenti +task: afrisenti_twi_prompt_3 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_3/afrisenti_yor.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_3/afrisenti_yor.yaml new file mode 100644 index 0000000000000000000000000000000000000000..9d1b7ac324b28b7a64220b6b318aabf9537594fc --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_3/afrisenti_yor.yaml @@ -0,0 +1,7 @@ +# Generated by utils.py +dataset_name: yor +doc_to_text: "You are an assistant able to detect sentiments in tweets. \n\nGiven\ + \ the sentiment labels Neutral, Positive or Negative; what is the sentiment of the\ + \ Yoruba statement below? Return only the labels. \n\ntext: {{tweet}} \nlabel:" +include: afrisenti +task: afrisenti_yor_prompt_3 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_3/utils.py b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_3/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..3e735e2deb1f9c53152c072615aebe8ba3acb90b --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_3/utils.py @@ -0,0 +1 @@ +from lm_eval.utils import weighted_f1_score diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_3/xx.py b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_3/xx.py new file mode 100644 index 0000000000000000000000000000000000000000..2133cfa0139de116c3d54e6c3866c5e4c26bbc53 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_3/xx.py @@ -0,0 +1,5 @@ +from datasets import load_dataset + + +data = load_dataset("masakhane/afrisenti", "por", trust_remote_code=True) +print(data) diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_4/afrisenti b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_4/afrisenti new file mode 100644 index 0000000000000000000000000000000000000000..6464d7b21693a1565f8479757a89a650cf84ff0c --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_4/afrisenti @@ -0,0 +1,39 @@ +tag: + - afrobench_sentiment_tasks + - afrisenti_prompt_4 +dataset_path: masakhane/afrisenti +dataset_name: null +dataset_kwargs: {trust_remote_code: True} +output_type: multiple_choice +validation_split: validation +test_split: test +fewshot_split: train +doc_to_target: label +doc_to_choice: + - "negative" + - "positive" + - "neutral" +should_decontaminate: true +doc_to_decontamination_query: 'text: {{tweet}} \nlabel: ' +metric_list: + - metric: f1 + aggregation: !function utils.weighted_f1_score + # aggregation: mean + average: weighted + hf_evaluate: true + higher_is_better: True + ignore_case: true + ignore_punctuation: true + regexes_to_ignore: + - "," + - "\\$" + - metric: acc + aggregation: mean + higher_is_better: true + ignore_case: true + ignore_punctuation: true + regexes_to_ignore: + - "," + - "\\$" +metadata: + version: 1.0 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_4/afrisenti_amh.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_4/afrisenti_amh.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8a30a72a45d38f6f4221a1e6ceaef93dd5472bbd --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_4/afrisenti_amh.yaml @@ -0,0 +1,6 @@ +# Generated by utils.py +dataset_name: amh +doc_to_text: "Label the following text as Neutral, Positive, or Negative. Provide\ + \ only the label as your response. \n\ntext: {{tweet}} \nlabel: " +include: afrisenti +task: afrisenti_amh_prompt_4 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_4/afrisenti_arq.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_4/afrisenti_arq.yaml new file mode 100644 index 0000000000000000000000000000000000000000..125771f5d6877585bb2b9a50121da7e5a56d7805 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_4/afrisenti_arq.yaml @@ -0,0 +1,6 @@ +# Generated by utils.py +dataset_name: arq +doc_to_text: "Label the following text as Neutral, Positive, or Negative. Provide\ + \ only the label as your response. \n\ntext: {{tweet}} \nlabel: " +include: afrisenti +task: afrisenti_arq_prompt_4 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_4/afrisenti_ary.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_4/afrisenti_ary.yaml new file mode 100644 index 0000000000000000000000000000000000000000..7868fbf3e6739cd69a4732a09b80cc08359ccbe8 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_4/afrisenti_ary.yaml @@ -0,0 +1,6 @@ +# Generated by utils.py +dataset_name: ary +doc_to_text: "Label the following text as Neutral, Positive, or Negative. Provide\ + \ only the label as your response. \n\ntext: {{tweet}} \nlabel: " +include: afrisenti +task: afrisenti_ary_prompt_4 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_4/afrisenti_hau.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_4/afrisenti_hau.yaml new file mode 100644 index 0000000000000000000000000000000000000000..5e7e9a443e5926dfe4de5074e1416417d38b4447 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_4/afrisenti_hau.yaml @@ -0,0 +1,6 @@ +# Generated by utils.py +dataset_name: hau +doc_to_text: "Label the following text as Neutral, Positive, or Negative. Provide\ + \ only the label as your response. \n\ntext: {{tweet}} \nlabel: " +include: afrisenti +task: afrisenti_hau_prompt_4 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_4/afrisenti_ibo.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_4/afrisenti_ibo.yaml new file mode 100644 index 0000000000000000000000000000000000000000..686e16c29e4aad0ff13e160befeb31e5b25a7f54 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_4/afrisenti_ibo.yaml @@ -0,0 +1,6 @@ +# Generated by utils.py +dataset_name: ibo +doc_to_text: "Label the following text as Neutral, Positive, or Negative. Provide\ + \ only the label as your response. \n\ntext: {{tweet}} \nlabel: " +include: afrisenti +task: afrisenti_ibo_prompt_4 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_4/afrisenti_por.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_4/afrisenti_por.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f5196bcf58303558d5a214caffde460b8675d08f --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_4/afrisenti_por.yaml @@ -0,0 +1,6 @@ +# Generated by utils.py +dataset_name: por +doc_to_text: "Label the following text as Neutral, Positive, or Negative. Provide\ + \ only the label as your response. \n\ntext: {{tweet}} \nlabel: " +include: afrisenti +task: afrisenti_por_prompt_4 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_5/afrisenti_amh.yaml b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_5/afrisenti_amh.yaml new file mode 100644 index 0000000000000000000000000000000000000000..866ffbe9fcab9e67a1ba4a9781dddd1ef60e8043 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/prompt_5/afrisenti_amh.yaml @@ -0,0 +1,13 @@ +# Generated by utils.py +dataset_name: amh +doc_to_text: "You are tasked with performing sentiment classification on the following\ + \ Amharic text. For each input, classify the sentiment as positive, negative, or\ + \ neutral. Use the following guidelines: \n\n Positive: The text expresses happiness,\ + \ satisfaction, or optimism. \nNegative: The text conveys disappointment, dissatisfaction,\ + \ or pessimism. \nNeutral: The text is factual, objective, or without strong emotional\ + \ undertones. \n\nIf the text contains both positive and negative sentiments, choose\ + \ the dominant sentiment. For ambiguous or unclear sentiments, select the label\ + \ that best reflects the overall tone. Please provide a single classification for\ + \ each input.\n\ntext: {{tweet}} \nlabel: " +include: afrisenti +task: afrisenti_amh_prompt_5 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/utils.py b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..b5f9b74e2eb12db6a985c8428830933f8adcc936 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrobench/afrisenti/utils.py @@ -0,0 +1,124 @@ +import argparse + +import yaml + + +class FunctionTag: + def __init__(self, value): + self.value = value + + +def prompt_func(mode, lang): + prompt_map = { + "prompt_1": "Does this statement; {{tweet}} have a Neutral, Positive or Negative sentiment? Labels only", + "prompt_2": f"Does this {lang} statement; " + "'{{tweet}}' have a Neutral, Positive or Negative sentiment? Labels only", + "prompt_3": f"You are an assistant able to detect sentiments in tweets. \n\n" + f"Given the sentiment labels Neutral, Positive or Negative; what is " + f"the sentiment of the {lang} statement below? Return only the labels. " + "\n\ntext: {{tweet}} \nlabel:", + "prompt_4": "Label the following text as Neutral, Positive, or Negative. Provide only the label as your " + "response. \n\ntext: {{tweet}} \nlabel: ", + "prompt_5": f"You are tasked with performing sentiment classification on the following {lang} text. " + f"For each input, classify the sentiment as positive, negative, or neutral. " + f"Use the following guidelines: \n\n " + f"Positive: The text expresses happiness, satisfaction, or optimism. \n" + f"Negative: The text conveys disappointment, dissatisfaction, or pessimism. \n" + f"Neutral: The text is factual, objective, or without strong emotional undertones. \n\n" + f"If the text contains both positive and negative sentiments, choose the dominant sentiment. " + f"For ambiguous or unclear sentiments, select the label that best reflects the overall tone. " + "Please provide a single classification for each input.\n\ntext: {{tweet}} \nlabel: ", + } + return prompt_map[mode] + + +def gen_lang_yamls(output_dir: str, overwrite: bool, mode: str) -> None: + """ + Generate a yaml file for each language. + + :param output_dir: The directory to output the files to. + :param overwrite: Whether to overwrite files if they already exist. + """ + err = [] + languages = { + "amh": "Amharic", + "arq": "Algerian Arabic", + "ary": "Moroccan Arabic", + "hau": "Hausa", + "ibo": "Igbo", + "kin": "Kinyarwanda", + "orm": "Oromo", + "pcm": "Nigerian Pidgin", + "por": "Mozambique Portuguese", + "swa": "Swahili", + "tir": "Tigrinya", + "tso": "Xithonga", + "twi": "Twi", + "yor": "Yoruba", + } + for lang in languages.keys(): + try: + file_name = f"afrisenti_{lang}.yaml" + task_name = f"afrisenti_{lang}_{mode}" + yaml_template = "afrisenti" + if int(mode.split("_")[-1]) > 1: + yaml_details = { + "include": yaml_template, + "task": task_name, + "dataset_name": lang, + "doc_to_text": prompt_func(mode, languages[lang]), + } + else: + yaml_details = { + "include": yaml_template, + "task": task_name, + "dataset_name": lang, + } + with open( + f"{output_dir}/{mode}/{file_name}", + "w" if overwrite else "x", + encoding="utf8", + ) as f: + f.write("# Generated by utils.py\n") + yaml.dump( + yaml_details, + f, + allow_unicode=True, + ) + except FileExistsError: + err.append(file_name) + + if len(err) > 0: + raise FileExistsError( + "Files were not created because they already exist (use --overwrite flag):" + f" {', '.join(err)}" + ) + + +def main() -> None: + """Parse CLI args and generate language-specific yaml files.""" + parser = argparse.ArgumentParser() + parser.add_argument( + "--overwrite", + default=True, + action="store_true", + help="Overwrite files if they already exist", + ) + parser.add_argument( + "--output-dir", + default="./", + help="Directory to write yaml files to", + ) + parser.add_argument( + "--mode", + default="prompt_1", + choices=["prompt_1", "prompt_2", "prompt_3", "prompt_4", "prompt_5"], + help="Prompt number", + ) + args = parser.parse_args() + + gen_lang_yamls(output_dir=args.output_dir, overwrite=args.overwrite, mode=args.mode) + + +if __name__ == "__main__": + main() diff --git a/obtain_metric.py b/obtain_metric.py new file mode 100644 index 0000000000000000000000000000000000000000..6c7119d74e62c0ddc60f65c86d815fe308580a66 --- /dev/null +++ b/obtain_metric.py @@ -0,0 +1,242 @@ +import json +import numpy as np +import os +from zscore import kurtosis_outlier_layers + +def mean_and_var(arr): + """ + 输入: arr (np.ndarray 或能转成 np.array 的对象) + 输出: (mean, variance) + """ + arr = np.asarray(arr, dtype=float) + mean = np.mean(arr) + var = np.var(arr) # 默认是总体方差，如果要无偏估计可用 ddof=1 + return mean, var + +def select_layers_shaped(kurtosis, alpha, k=5, return_scores=False, + pos_gaussians=None, neg_gaussians=None): + """ + 强塑形联合指标: + score_i = (alpha_i / kurtosis_i) * ( 1 + + sum_j b_j * exp(-((kurtosis_i - c_j)^2) / s_j^2) + - sum_t d_t * exp(-((kurtosis_i - u_t)^2) / r_t^2) ) + + 参数 + ---- + kurtosis : list/ndarray κ + alpha : list/ndarray α + k : 取前 k 个下标 + pos_gaussians : [(c, s, b), ...] 正向高斯(中心c, 带宽s>0, 系数b>0) + neg_gaussians : [(u, r, d), ...] 负向高斯(中心u, 带宽r>0, 系数d>0) + + 返回 + ---- + idx_topk : list[int] + (可选) scores : ndarray + """ + krt = np.asarray(kurtosis, dtype=float) + alp = np.asarray(alpha, dtype=float) + if krt.shape != alp.shape: + raise ValueError("kurtosis 与 alpha 形状不一致") + + base = alp - krt + + c, sigma = mean_and_var(krt) + beta = 1.5 + weight = 1.0 + beta * np.exp(-((krt - c) ** 2) / (sigma ** 2)) + scores = base * weight + return scores + +def joint_score(kurtosis, alpha, + alpha0=4.20, s_alpha=0.090, + s_k=1.254, w_alpha=9.48, w_k=1.02): + """ + Kurtosis-Alpha 联合指标 (Taguchi 损失型) + ------------------------- + 参数 + kurtosis : list/ndarray + 每层的 kurtosis 值 + alpha : list/ndarray + 每层的 alpha 值 + alpha0 : float + alpha 的目标值（名义最佳点） + s_alpha : float + alpha 的尺度因子 + s_k : float + kurtosis 的尺度因子 + w_alpha, w_k : float + alpha 和 kurtosis 的权重 + + 返回 + scores : ndarray + 每层的综合得分（越大越好） + """ + k = np.array(kurtosis) + a = np.array(alpha) + + loss_k = (k / s_k) ** 2 + loss_a = ((a - alpha0) / s_alpha) ** 2 + + scores = -(w_k * loss_k + w_alpha * loss_a) + return scores + +def simple_joint_score(kurtosis, alpha, stable_rank): + """ + 极简联合指标: + S_l = -k_l * (alpha_l - Q80(alpha))^2 + + 参数 + ---- + kurtosis : list/ndarray + 每层的 kurtosis 值 + alpha : list/ndarray + 每层的 alpha 值 + + 返回 + ---- + scores : ndarray + 每层的综合得分（越大越好） + """ + k = np.array(kurtosis) + a = np.array(alpha) + s = np.array(stable_rank) + print('alpha_hat_datas') + # h = np.array(alpha_hat_datas) + k = (k - k.min()) / (k.max() - k.min()) + # k_exp = np.exp(k) # 防止溢出 + # k = k_exp / np.sum(k_exp) + + a = 1 / a + a = (a - a.min()) / (a.max() - a.min()) + # a_exp = np.exp(a) # 防止溢出 + # a = a_exp / np.sum(a_exp) + + s = 1 / s + s = (s - s.min()) / (s.max() - s.min()) + # s_exp = np.exp(s) + # s = s_exp / np.sum(s_exp) + + # h = 1 / h + # # h= (h - h.min()) / (h.max() - h.min()) + # h_exp = np.exp(h) # 防止溢出 + # h = h_exp / np.sum(h_exp) + + # alpha_target = np.percentile(a, 80) # α 的 80 分位数 + # # alpha_target = 0 + # scores = -k * (a - alpha_target) ** 2 + + # k = k * 0.5 + # h = h * 0.5 + # print(k) + # print(a) + # print(h) + # a = a + 0.3 * h + + + # a = a +1 + # k = k +1 + # h = h +1 + print('a', {i:aa for i, aa in enumerate(a)}) + print() + print('k', {i:aa for i, aa in enumerate(k)}) + print() + print('s', {i:aa for i, aa in enumerate(s)}) + # k = 100 * k + # scores = (k / a) * (k - a) + k = 0.2 * k + s = 1.2 * s + scores = k + a + s + # scores = (k * h) * (k + h) + # print() + # print('scores', {i:aa for i, aa in enumerate(scores)}) + return scores + +model_names = ['Llama-2-7b-hf',] + # 'Llama-2-13b-hf', 'Qwen3-8B', 'Qwen3-4B', 'Mistral-7B-Instruct-v0.3','Llama-3.2-3B-Instruct'] + +for model in model_names: + # model = "Llama-2-7b-hf" + print(model) + alpha_path = f'/mnt/bn/life-mllm/users/cxr/quantization/quantization_metric/alpha_values/alpha_values_{model}.json' + kurtosis_path = f'/mnt/bn/life-mllm/users/cxr/quantization/lm-quant-toolkit/kurtosis_means/kurtosis_means-{model}.json' + alpha_hat_path = f'/mnt/bn/life-mllm/users/cxr/quantization/quantization_metric/alpha_hat/alpha_values_{model}.json' + stable_rank_path = f'/mnt/bn/life-mllm/users/cxr/quantization/quantization_metric/stable_rank/stable_rank_{model}.json' + zd_path = f'/mnt/bn/life-mllm/users/cxr/quantization/quantization_metric/ZD/ZD_{model}.json' + bi_path = f'/mnt/bn/life-mllm/users/cxr/quantization/quantization_metric/BI/BI_{model}.json' + + with open(stable_rank_path, 'r', encoding='utf-8') as f: + stable_rank_datas = json.load(f) + + with open(alpha_path, 'r', encoding='utf-8') as f: + alpha_datas = json.load(f) + + with open(kurtosis_path, 'r', encoding='utf-8') as f: + kurtosis_datas = json.load(f) + + with open(alpha_hat_path, 'r', encoding='utf-8') as f: + alpha_hat_datas = json.load(f) + + with open(zd_path, 'r', encoding='utf-8') as f: + zd_datas = json.load(f) + + with open(bi_path, 'r', encoding='utf-8') as f: + bi_datas = json.load(f) + + print(len(kurtosis_datas)) + print('kurtosis_datas', np.argsort(kurtosis_datas)) + print('stable_rank_datas', np.argsort([-a for a in stable_rank_datas])) + print('alpha_datas', np.argsort([-a for a in alpha_datas])) + print('alpha_hat_datas', np.argsort([-a for a in alpha_hat_datas])) + + print('zd_datas', np.argsort([-a for a in zd_datas])) + print('bi_datas', np.argsort([a for a in bi_datas])) + + scores = simple_joint_score(kurtosis_datas, alpha_datas, stable_rank_datas) + idx_sorted = np.argsort(scores) + # print('scores:', scores) + print('idx_sorted:',idx_sorted) + print() + + kurtosis_idx = np.argsort(kurtosis_datas).tolist() + stable_rank_idx = np.argsort([-a for a in stable_rank_datas]).tolist() + alpha_idx = np.argsort([-a for a in alpha_datas]).tolist() + bi_idx = np.argsort([a for a in bi_datas]).tolist() + z_idx = kurtosis_outlier_layers(kurtosis_datas) + zd_idx = np.argsort([-a for a in zd_datas]).tolist() + layrs = [5, 10] + + for layr in layrs: + # bits = [4 for _ in range(len(alpha_datas))] + # print(kurtosis_idx[:layr]) + # for i in kurtosis_idx[:layr]: + # bits[i] = 2 + # with open(os.path.join('/mnt/bn/life-mllm/users/cxr/quantization/baselines', f'{model}_kurtosis_idx_{layr}.json'), "w", encoding="utf-8") as f: + # json.dump(bits, f, ensure_ascii=False, indent=4) + + # bits = [4 for _ in range(len(alpha_datas))] + # for i in alpha_idx[:layr]: + # bits[i] = 2 + # with open(os.path.join('/mnt/bn/life-mllm/users/cxr/quantization/baselines', f'{model}_alpha_idx_{layr}.json'), "w", encoding="utf-8") as f: + # json.dump(bits, f, ensure_ascii=False, indent=4) + + # bits = [4 for _ in range(len(alpha_datas))] + # for i in z_idx[:layr]: + # bits[i] = 2 + # with open(os.path.join('/mnt/bn/life-mllm/users/cxr/quantization/baselines', f'{model}_z_idx_{layr}.json'), "w", encoding="utf-8") as f: + # json.dump(bits, f, ensure_ascii=False, indent=4) + + # bits = [4 for _ in range(len(alpha_datas))] + # print(bi_idx[:layr]) + # for i in bi_idx[:layr]: + # bits[i] = 2 + # with open(os.path.join('/mnt/bn/life-mllm/users/cxr/quantization/baselines_1', f'{model}_bi_idx_{layr}.json'), "w", encoding="utf-8") as f: + # json.dump(bits, f, ensure_ascii=False, indent=4) + + kurtosis_datas, stable_rank_datas + bits = [4 for _ in range(len(alpha_datas))] + print(idx_sorted[:layr]) + for i in idx_sorted[:layr]: + bits[i] = 2 + with open(os.path.join('/mnt/bn/life-mllm/users/cxr/quantization/ours2', f'{model}_aks_plus_idx_sorted_{layr}.json'), "w", encoding="utf-8") as f: + json.dump(bits, f, ensure_ascii=False, indent=4) + \ No newline at end of file diff --git a/test.py b/test.py new file mode 100644 index 0000000000000000000000000000000000000000..c89a5ffe659a7f7b9d6d95451eda87c8dfed62fc --- /dev/null +++ b/test.py @@ -0,0 +1,24 @@ +from huggingface_hub import HfApi + +# 初始化 API +api = HfApi() + +# 上传超大文件夹（自动分片+续传） +api.upload_large_folder( + folder_path="./clean-model-files", # 清理后的本地目录 + repo_id="你的用户名/你的仓库名", # 如 zhangsan/my-llm-model + repo_type="model", # 仓库类型 + allow_patterns=[ # 仅上传以下文件（白名单） + "*.safetensors", + "*.bin", + "config.json", + "tokenizer*.json", + "vocab.txt" + ], + ignore_patterns=[ # 排除以下文件（黑名单） + ".cache/**", + "*.log", + "__pycache__/**" + ], + overwrite=True, # 覆盖已有文件 +) \ No newline at end of file diff --git a/test.sh b/test.sh new file mode 100644 index 0000000000000000000000000000000000000000..b1d9ab1fcadb27239fb01db109f84369766da043 --- /dev/null +++ b/test.sh @@ -0,0 +1 @@ +conda activate lm-eval \ No newline at end of file diff --git a/test2.py b/test2.py new file mode 100644 index 0000000000000000000000000000000000000000..6af55604589ae5b7a5f9f1f2541a588eee865aff --- /dev/null +++ b/test2.py @@ -0,0 +1,51 @@ +import json +import heapq +def get_top_k_indices(json_file_path, k): + """ + 读取JSON文件中的列表，返回最大的k个元素的索引 + + Args: + json_file_path: JSON文件路径 + k: 需要获取的最大元素的个数 + + Returns: + list: 按元素大小降序排列的索引列表 + + Raises: + FileNotFoundError: 文件不存在时抛出 + ValueError: k值无效或数据格式错误时抛出 + """ + # 读取JSON文件 + try: + with open(json_file_path, 'r', encoding='utf-8') as f: + data = json.load(f) + except FileNotFoundError: + raise FileNotFoundError(f"文件 {json_file_path} 不存在") + except json.JSONDecodeError: + raise ValueError("JSON文件格式错误") + + # 验证数据类型 + if not isinstance(data, list): + raise ValueError("JSON文件内容不是列表") + + # 验证k值的有效性 + if k <= 0 or k > len(data): + raise ValueError(f"k值无效，应在1到{len(data)}之间") + + # 获取元素值和索引的元组列表 [(value, index), ...] + value_index_pairs = [(value, idx) for idx, value in enumerate(data)] + + # 方法1：使用heapq获取最大的k个元素（效率更高，O(n log k)） + top_k_pairs = heapq.nlargest(k, value_index_pairs, key=lambda x: x[0]) + + # 方法2：使用排序（简单直观，O(n log n)） + # sorted_pairs = sorted(value_index_pairs, key=lambda x: x[0], reverse=True) + # top_k_pairs = sorted_pairs[:k] + + # 提取索引 + top_k_indices = [pair[1] for pair in top_k_pairs] + + return top_k_indices + +a = get_top_k_indices('/mnt/bn/life-mllm/users/cxr/quantization/quantization_metric/metrics/alpha/alpha_mlp_Llama-2-7b-hf.json', 10) +print(a) \ No newline at end of file diff --git a/zscore.py b/zscore.py new file mode 100644 index 0000000000000000000000000000000000000000..45bc481e0d11acdeaea0e782520b02bea5dc6e05 --- /dev/null +++ b/zscore.py @@ -0,0 +1,44 @@ +import numpy as np +import json + +def kurtosis_outlier_layers(kurtosis_values, threshold=3.0, method="subtract"): + """ + 根据 Kurtosis 值列表，计算差分并用 z-score 检测异常层 + + 参数: + kurtosis_values (list or np.ndarray): Kurtosis 值序列 (s1, s2, ..., sn) + threshold (float): z-score 阈值 (默认 3.0) + method (str): 差分方式 + "subtract" -> si+1 - si + "divide" -> si+1 / si + + 返回: + z_scores (np.ndarray): 差分的 z-score + outlier_indices (list): 被判定为异常的层索引 (对应原始 Kurtosis 序列中的层号) + """ + values = np.array(kurtosis_values, dtype=float) + + # Step 1: 差分 + if method == "subtract": + diffs = np.diff(values) # s2-s1, s3-s2, ... + elif method == "divide": + diffs = values[1:] / values[:-1] + else: + raise ValueError("method must be 'subtract' or 'divide'") + + # Step 2: 计算 z-score + mean = np.mean(diffs) + std = np.std(diffs, ddof=1) # 样本标准差 + z_scores = abs((diffs - mean) / std) + results = [(i+1, z) for i, z in enumerate(z_scores)] + results.sort(key=lambda x: x[1], reverse=False) + print(results) + results = [i[0] for i in results] + + return results + +k_path = '/mnt/bn/life-mllm/users/cxr/quantization/lm-quant-toolkit/kurtosis_means/kurtosis_means-Llama-2-7b-hf.json' +with open(k_path, "r", encoding="utf-8") as f: + kurtosis_values = json.load(f) + +kurtosis_outlier_layers(kurtosis_values) \ No newline at end of file