From f20570343f7895c61d814bf53bcbb41bb0fb5f22 Mon Sep 17 00:00:00 2001 From: Andrey Glushenkov Date: Thu, 6 Apr 2023 02:29:36 +0300 Subject: [PATCH] GPTQv2 support GPTQv2 support. 1. Adds dependency on `triton` 2. Refactors autograd_4bit to include both GPTQv1 and GPTQv2 3. Introduces new environment variable GPTQ_VERSION to select autograd_4bit version 4. Fixes triton kernels 5. Matrix multiplications are in fp16 --- Finetune4bConfig.py | 3 +- README.md | 10 ++ autograd_4bit/__init__.py | 21 +++ autograd_4bit/autograd_4bit_v1.py | 208 ++++++++++++++++++++++++++++ autograd_4bit/autograd_4bit_v2.py | 221 ++++++++++++++++++++++++++++++ finetune.py | 1 + requirements.txt | 1 + triton_test.py | 154 +++++++++++++++++++++ 8 files changed, 618 insertions(+), 1 deletion(-) create mode 100644 autograd_4bit/__init__.py create mode 100644 autograd_4bit/autograd_4bit_v1.py create mode 100644 autograd_4bit/autograd_4bit_v2.py create mode 100644 triton_test.py diff --git a/Finetune4bConfig.py b/Finetune4bConfig.py index a8a33bf..3880462 100644 --- a/Finetune4bConfig.py +++ b/Finetune4bConfig.py @@ -95,5 +95,6 @@ class Finetune4bConfig: f"{self.warmup_steps=}\n{self.save_steps=}\n{self.save_total_limit=}\n" +\ f"{self.logging_steps=}\n" +\ f"{self.checkpoint=}\n{self.skip=}\n" +\ - f"{self.world_size=}\n{self.ddp=}\n{self.device_map=}" + f"{self.world_size=}\n{self.ddp=}\n{self.device_map=}\n" +\ + f"{self.groupsize=}\n" return s.replace("self.", "") diff --git a/README.md b/README.md index 2326cde..161bb26 100644 --- a/README.md +++ b/README.md @@ -34,10 +34,20 @@ pip install -r requirements.txt ~The same finetune script from https://github.com/tloen/alpaca-lora can be used.~
After installation, this script can be used: +GPTQv1: ``` python finetune.py ``` +or +``` +GPTQ_VERSION=1 python finetune.py +``` + +GPTQv2: +``` +GPTQ_VERSION=2 python finetune.py +``` # Inference diff --git a/autograd_4bit/__init__.py b/autograd_4bit/__init__.py new file mode 100644 index 0000000..bee84b3 --- /dev/null +++ b/autograd_4bit/__init__.py @@ -0,0 +1,21 @@ +import os +from colorama import init, Fore, Back, Style +init(autoreset=True) + +try: + GPTQ_VERSION = int(os.environ["GPTQ_VERSION"]) +except: + print(Style.BRIGHT + Fore.YELLOW + "GPTQ_VERSION environment not provided. Fallback to GPTQv1") + GPTQ_VERSION = 1 # Fallback version + +loader = None + + +if GPTQ_VERSION == 1: + from .autograd_4bit_v1 import Autograd4bitQuantLinear, load_llama_model_4bit_low_ram + print(Style.BRIGHT + Fore.GREEN + "GPTQv1 set") +elif GPTQ_VERSION == 2: + from .autograd_4bit_v2 import Autograd4bitQuantLinear, load_llama_model_4bit_low_ram + print(Style.BRIGHT + Fore.GREEN + "GPTQv2 set") +else: + raise ValueError("GPTQ_VERSION not set or invalid") \ No newline at end of file diff --git a/autograd_4bit/autograd_4bit_v1.py b/autograd_4bit/autograd_4bit_v1.py new file mode 100644 index 0000000..cb9d308 --- /dev/null +++ b/autograd_4bit/autograd_4bit_v1.py @@ -0,0 +1,208 @@ +import matmul_utils_4bit as mm4b +import torch +import torch.nn as nn +import time + + +class AutogradMatmul4bit(torch.autograd.Function): + + @staticmethod + def forward(ctx, x, qweight, scales, zeros, groupsize=-1): + ctx.save_for_backward(qweight, scales, zeros) + ctx.groupsize = groupsize + if groupsize == -1: + output = mm4b._matmul4bit_v1_recons(x, qweight, scales, zeros) + else: + output = mm4b._matmul4bit_v2_recons(x, qweight, scales, zeros, groupsize) + output = output.clone() + return output + + @staticmethod + def backward(ctx, grad_output): + qweight, scales, zeros = ctx.saved_tensors + groupsize = ctx.groupsize + if groupsize == -1: + grad = mm4b._matmul4bit_v1_recons(grad_output, qweight, scales, zeros, transpose=True) + else: + grad = mm4b._matmul4bit_v2_recons(grad_output, qweight, scales, zeros, groupsize=groupsize, transpose=True) + return grad, None, None, None, None + + +# Assumes layer is perfectly divisible into 256 * 256 blocks +class Autograd4bitQuantLinear(nn.Module): + + def __init__(self, in_features, out_features, groupsize=None): + super().__init__() + bits = 4 + self.in_features = in_features + self.out_features = out_features + self.bits = bits + self.register_buffer('zeros', torch.empty((out_features, 1))) + self.register_buffer('scales', torch.empty((out_features, 1))) + self.bias = nn.Parameter(torch.empty(out_features)) + self.register_buffer( + 'qweight', torch.empty((in_features // 256 * (bits * 8), out_features), dtype=torch.int) + ) + + + def forward(self, x): + if torch.is_grad_enabled(): + out = AutogradMatmul4bit.apply(x, self.qweight, self.scales, + self.qzeros if self.groupsize != -1 else self.zeros, self.groupsize) + out.add_(self.bias) + else: + out = mm4b.matmul4bit(x, self.qweight, self.scales, + self.qzeros if self.groupsize != -1 else self.zeros, self.groupsize) + out.add_(self.bias) + return out + + +def make_quant_for_4bit_autograd(module, names, name='', groupsize=-1): + if isinstance(module, Autograd4bitQuantLinear): + return + for attr in dir(module): + tmp = getattr(module, attr) + name1 = name + '.' + attr if name != '' else attr + if name1 in names: + setattr( + module, attr, Autograd4bitQuantLinear(tmp.in_features, tmp.out_features, groupsize=groupsize) + ) + for name1, child in module.named_children(): + make_quant_for_4bit_autograd(child, names, name + '.' + name1 if name != '' else name1, groupsize=groupsize) + + +def model_to_half(model): + model.half() + for n, m in model.named_modules(): + if isinstance(m, Autograd4bitQuantLinear): + m.zeros = m.zeros.half() + m.scales = m.scales.half() + m.bias = m.bias.half() + print('Converted as Half.') + + +def model_to_float(model): + model.float() + for n, m in model.named_modules(): + if isinstance(m, Autograd4bitQuantLinear): + m.zeros = m.zeros.float() + m.scales = m.scales.float() + m.bias = m.bias.float() + print('Converted as Float.') + + +def find_layers(module, layers=[nn.Conv2d, nn.Linear], name=''): + if type(module) in layers: + return {name: module} + res = {} + for name1, child in module.named_children(): + res.update(find_layers( + child, layers=layers, name=name + '.' + name1 if name != '' else name1 + )) + return res + + +def load_llama_model_4bit_low_ram(config_path, model_path, groupsize=-1, half=False, device_map="auto", seqlen=2048): + import accelerate + from transformers import LlamaConfig, LlamaForCausalLM, LlamaTokenizer + + print("Loading Model ...") + t0 = time.time() + + with accelerate.init_empty_weights(): + config = LlamaConfig.from_pretrained(config_path) + model = LlamaForCausalLM(config) + model = model.eval() + layers = find_layers(model) + for name in ['lm_head']: + if name in layers: + del layers[name] + make_quant_for_4bit_autograd(model, layers, groupsize=groupsize) + model = accelerate.load_checkpoint_and_dispatch( + model=model, + checkpoint=model_path, + device_map=device_map, + no_split_module_classes=["LlamaDecoderLayer"] + ) + + model.seqlen = seqlen + + if half: + model_to_half(model) + + tokenizer = LlamaTokenizer.from_pretrained(config_path) + tokenizer.truncation_side = 'left' + + print(f"Loaded the model in {(time.time()-t0):.2f} seconds.") + + return model, tokenizer + +def load_llama_model_4bit_low_ram_and_offload_to_cpu(config_path, model_path, lora_path=None, groupsize=-1, seqlen=2048, max_memory=None): + import accelerate + from transformers import LlamaConfig, LlamaForCausalLM, LlamaTokenizer + + if max_memory is None: + max_memory = {0: '24Gib', 'cpu': '48Gib'} + + print("Loading Model ...") + t0 = time.time() + + with accelerate.init_empty_weights(): + config = LlamaConfig.from_pretrained(config_path) + model = LlamaForCausalLM(config) + model = model.eval() + layers = find_layers(model) + for name in ['lm_head']: + if name in layers: + del layers[name] + make_quant_for_4bit_autograd(model, layers, groupsize=groupsize) + accelerate.load_checkpoint_in_model(model, checkpoint=model_path, device_map={'': 'cpu'}) + + # rotary_emb fix + for n, m in model.named_modules(): + if 'rotary_emb' in n: + cos_cached = m.cos_cached.clone().cpu() + sin_cached = m.sin_cached.clone().cpu() + break + + if lora_path is not None: + from peft import PeftModel + from peft.tuners.lora import Linear4bitLt + model = PeftModel.from_pretrained(model, lora_path, device_map={'': 'cpu'}, torch_dtype=torch.float32) + print('{} Lora Applied.'.format(lora_path)) + + model.seqlen = seqlen + + print('Apply half ...') + for n, m in model.named_modules(): + if isinstance(m, Autograd4bitQuantLinear) or ((lora_path is not None) and isinstance(m, Linear4bitLt)): + m.zeros = m.zeros.half() + m.scales = m.scales.half() + m.bias = m.bias.half() + + print('Dispatching model ...') + device_map = accelerate.infer_auto_device_map(model, max_memory=max_memory, no_split_module_classes=["LlamaDecoderLayer"]) + model = accelerate.dispatch_model(model, device_map=device_map, offload_buffers=True, main_device=0) + torch.cuda.empty_cache() + print('Total {:.2f} Gib VRAM used.'.format(torch.cuda.memory_allocated() / 1024 / 1024)) + + # rotary_emb fix + for n, m in model.named_modules(): + if 'rotary_emb' in n: + if getattr(m, '_hf_hook', None): + if isinstance(m._hf_hook, accelerate.hooks.SequentialHook): + hooks = m._hf_hook.hooks + else: + hooks = [m._hf_hook] + for hook in hooks: + if hook.offload: + if n + '.sin_cached' not in hook.weights_map.dataset.state_dict.keys(): + hook.weights_map.dataset.state_dict[n + '.sin_cached'] = sin_cached.clone().cpu() + hook.weights_map.dataset.state_dict[n + '.cos_cached'] = cos_cached.clone().cpu() + + tokenizer = LlamaTokenizer.from_pretrained(config_path) + tokenizer.truncation_side = 'left' + + print(f"Loaded the model in {(time.time()-t0):.2f} seconds.") + + return model, tokenizer diff --git a/autograd_4bit/autograd_4bit_v2.py b/autograd_4bit/autograd_4bit_v2.py new file mode 100644 index 0000000..20c253d --- /dev/null +++ b/autograd_4bit/autograd_4bit_v2.py @@ -0,0 +1,221 @@ +from colorama import init, Fore, Back, Style +import torch +import torch.nn as nn +import time +import math +import triton +from triton_utils import matmul_248_kernel, trans_matmul_248_kernel + + +class AutogradMatmul4bit(torch.autograd.Function): + @staticmethod + def forward(ctx, input, qweight, scales, qzeros, g_idx, bits, maxq): + output = torch.empty((input.shape[0], qweight.shape[1]), device='cuda', dtype=torch.float16) + grid = lambda META: (triton.cdiv(input.shape[0], META['BLOCK_SIZE_M']) * triton.cdiv(qweight.shape[1], META['BLOCK_SIZE_N']),) + matmul_248_kernel[grid](input, qweight, output, + scales, qzeros, g_idx, + input.shape[0], qweight.shape[1], input.shape[1], bits, maxq, + input.stride(0), input.stride(1), + qweight.stride(0), qweight.stride(1), + output.stride(0), output.stride(1), + scales.stride(0), qzeros.stride(0)) + + ctx.save_for_backward(qweight, scales, qzeros, g_idx) + ctx.input_shape, ctx.bits,ctx.maxq = input.shape,bits, maxq + return output + + @staticmethod + def backward(ctx, grad_output): + input_shape, bits, maxq = ctx.input_shape, ctx.bits, ctx.maxq + qweight, scales, qzeros, g_idx = ctx.saved_tensors + grade_input = None + + if ctx.needs_input_grad[0]: + grade_input = torch.empty((input_shape[0], input_shape[1]), device='cuda', dtype=torch.float32) + grid = lambda META: (triton.cdiv(input_shape[0], META['BLOCK_SIZE_M']) * triton.cdiv(input_shape[1], META['BLOCK_SIZE_K']),) + trans_matmul_248_kernel[grid](grad_output, qweight, grade_input, + scales, qzeros, g_idx, + input_shape[0], qweight.shape[1], input_shape[1], bits, maxq, + grad_output.stride(0), grad_output.stride(1), + qweight.stride(0), qweight.stride(1), + grade_input.stride(0), grade_input.stride(1), + scales.stride(0), qzeros.stride(0)) + return grade_input, None, None, None, None, None, None + + +class Autograd4bitQuantLinear(nn.Module): + + def __init__(self, in_features, out_features, groupsize, bias=True): + super().__init__() + self.in_features = in_features + self.out_features = out_features + self.bits = 4 # Hardcoded 4-bits quantizations + self.maxq = 2 ** self.bits - 1 + self.groupsize = groupsize if groupsize != -1 else in_features + + self.register_buffer('qweight', torch.zeros((in_features // 32 * self.bits, out_features), dtype=torch.int32)) + self.register_buffer('qzeros', torch.zeros((math.ceil(in_features / self.groupsize), out_features // 32 * self.bits), dtype=torch.int32)) + self.register_buffer('scales', torch.zeros((math.ceil(in_features / self.groupsize), out_features), dtype=torch.float16)) + self.register_buffer('g_idx', torch.tensor([i // self.groupsize for i in range(in_features)], dtype = torch.int32)) + if bias: + self.register_buffer('bias', torch.zeros(out_features,dtype=torch.float16)) + else: + self.bias = None + + def forward(self, x): + out_shape = x.shape[:-1] + (self.out_features, ) + out = AutogradMatmul4bit.apply(x.reshape(-1,x.shape[-1]), self.qweight, self.scales, + self.qzeros, self.g_idx, self.bits, self.maxq) + out = out + self.bias if self.bias is not None else out + return out.reshape(out_shape) + + +def make_quant_for_4bit_autograd(module, names, name='', groupsize=-1): + if isinstance(module, Autograd4bitQuantLinear): + return + for attr in dir(module): + tmp = getattr(module, attr) + name1 = name + '.' + attr if name != '' else attr + if name1 in names: + setattr( + module, attr, Autograd4bitQuantLinear(tmp.in_features, tmp.out_features, groupsize=groupsize) + ) + for name1, child in module.named_children(): + make_quant_for_4bit_autograd(child, names, name + '.' + name1 if name != '' else name1, groupsize=groupsize) + + +def model_to_half(model): + model.half() + for n, m in model.named_modules(): + if isinstance(m, Autograd4bitQuantLinear): + m.qzeros = m.qzeros.half() + m.scales = m.scales.half() + m.bias = m.bias.half() + print(Style.BRIGHT + Fore.YELLOW + 'Converted as Half.') + + +def model_to_float(model): + model.float() + for n, m in model.named_modules(): + if isinstance(m, Autograd4bitQuantLinear): + m.qzeros = m.qzeros.float() + m.scales = m.scales.float() + m.bias = m.bias.float() + print(Style.BRIGHT + Fore.YELLOW + 'Converted as Float.') + + +def find_layers(module, layers=[nn.Conv2d, nn.Linear], name=''): + if type(module) in layers: + return {name: module} + res = {} + for name1, child in module.named_children(): + res.update(find_layers( + child, layers=layers, name=name + '.' + name1 if name != '' else name1 + )) + return res + + +def load_llama_model_4bit_low_ram(config_path, model_path, groupsize=-1, half=False, device_map="auto", seqlen=2048): + import accelerate + from transformers import LlamaConfig, LlamaForCausalLM, LlamaTokenizer + + print(Style.BRIGHT + Fore.CYAN + "Loading Model ...") + t0 = time.time() + + with accelerate.init_empty_weights(): + config = LlamaConfig.from_pretrained(config_path) + model = LlamaForCausalLM(config) + model = model.eval() + layers = find_layers(model) + for name in ['lm_head']: + if name in layers: + del layers[name] + make_quant_for_4bit_autograd(model, layers, groupsize=groupsize) + model = accelerate.load_checkpoint_and_dispatch( + model=model, + checkpoint=model_path, + device_map=device_map, + no_split_module_classes=["LlamaDecoderLayer"] + ) + + model.seqlen = seqlen + + if half: + model_to_half(model) + + tokenizer = LlamaTokenizer.from_pretrained(config_path) + tokenizer.truncation_side = 'left' + + print(Style.BRIGHT + Fore.GREEN + f"Loaded the model in {(time.time()-t0):.2f} seconds.") + + return model, tokenizer + +def load_llama_model_4bit_low_ram_and_offload_to_cpu(config_path, model_path, lora_path=None, groupsize=-1, seqlen=2048, max_memory=None): + import accelerate + from transformers import LlamaConfig, LlamaForCausalLM, LlamaTokenizer + + if max_memory is None: + max_memory = {0: '24Gib', 'cpu': '48Gib'} + + print(Style.BRIGHT + Fore.CYAN + "Loading Model ...") + t0 = time.time() + + with accelerate.init_empty_weights(): + config = LlamaConfig.from_pretrained(config_path) + model = LlamaForCausalLM(config) + model = model.eval() + layers = find_layers(model) + for name in ['lm_head']: + if name in layers: + del layers[name] + make_quant_for_4bit_autograd(model, layers, groupsize=groupsize) + accelerate.load_checkpoint_in_model(model, checkpoint=model_path, device_map={'': 'cpu'}) + + # rotary_emb fix + for n, m in model.named_modules(): + if 'rotary_emb' in n: + cos_cached = m.cos_cached.clone().cpu() + sin_cached = m.sin_cached.clone().cpu() + break + + if lora_path is not None: + from peft import PeftModel + from peft.tuners.lora import Linear4bitLt + model = PeftModel.from_pretrained(model, lora_path, device_map={'': 'cpu'}, torch_dtype=torch.float32) + print(Style.BRIGHT + Fore.GREEN + '{} Lora Applied.'.format(lora_path)) + + model.seqlen = seqlen + + print('Apply half ...') + for n, m in model.named_modules(): + if isinstance(m, Autograd4bitQuantLinear) or ((lora_path is not None) and isinstance(m, Linear4bitLt)): + m.qzeros = m.qzeros.half() + m.scales = m.scales.half() + m.bias = m.bias.half() + + print('Dispatching model ...') + device_map = accelerate.infer_auto_device_map(model, max_memory=max_memory, no_split_module_classes=["LlamaDecoderLayer"]) + model = accelerate.dispatch_model(model, device_map=device_map, offload_buffers=True, main_device=0) + torch.cuda.empty_cache() + print(Style.BRIGHT + Fore.YELLOW + 'Total {:.2f} Gib VRAM used.'.format(torch.cuda.memory_allocated() / 1024 / 1024)) + + # rotary_emb fix + for n, m in model.named_modules(): + if 'rotary_emb' in n: + if getattr(m, '_hf_hook', None): + if isinstance(m._hf_hook, accelerate.hooks.SequentialHook): + hooks = m._hf_hook.hooks + else: + hooks = [m._hf_hook] + for hook in hooks: + if hook.offload: + if n + '.sin_cached' not in hook.weights_map.dataset.state_dict.keys(): + hook.weights_map.dataset.state_dict[n + '.sin_cached'] = sin_cached.clone().cpu() + hook.weights_map.dataset.state_dict[n + '.cos_cached'] = cos_cached.clone().cpu() + + tokenizer = LlamaTokenizer.from_pretrained(config_path) + tokenizer.truncation_side = 'left' + + print(Style.BRIGHT + Fore.GREEN + f"Loaded the model in {(time.time()-t0):.2f} seconds.") + + return model, tokenizer diff --git a/finetune.py b/finetune.py index f374e2b..80d2156 100644 --- a/finetune.py +++ b/finetune.py @@ -109,6 +109,7 @@ if not ft_config.skip: per_device_train_batch_size=ft_config.mbatch_size, gradient_accumulation_steps=ft_config.gradient_accumulation_steps, warmup_steps=ft_config.warmup_steps, + optim="adamw_torch", num_train_epochs=ft_config.epochs, learning_rate=ft_config.lr, fp16=True, diff --git a/requirements.txt b/requirements.txt index 605c0d1..9b117d7 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,6 +4,7 @@ bitsandbytes datasets sentencepiece safetensors +triton git+https://github.com/huggingface/transformers.git git+https://github.com/sterlind/GPTQ-for-LLaMa.git@lora_4bit git+https://github.com/sterlind/peft.git diff --git a/triton_test.py b/triton_test.py new file mode 100644 index 0000000..eeb77a9 --- /dev/null +++ b/triton_test.py @@ -0,0 +1,154 @@ +import torch + +import triton +import triton.language as tl + +# % +# :code:`triton.jit`'ed functions can be auto-tuned by using the `triton.autotune` +# decorator, which consumes: +# - A list of :code:`triton.Config` objects that define different configurations of +# meta-parameters (e.g., BLOCK_SIZE_M) and compilation options (e.g., num_warps) to try +# - An autotuning *key* whose change in values will trigger evaluation of all the +# provided configs + + +@triton.autotune( + configs=[ + triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=3, num_warps=8), + triton.Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=3, num_warps=8), + triton.Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4), + triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4), + triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4), + triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4), + triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4), + triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4), + triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=5, num_warps=2), + triton.Config({'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=5, num_warps=2), + ], + key=['M', 'N', 'K'], +) +@triton.jit +def matmul_kernel( + # Pointers to matrices + a_ptr, b_ptr, c_ptr, + # Matrix dimensions + M, N, K, + # The stride variables represent how much to increase the ptr by when moving by 1 + # element in a particular dimension. E.g. stride_am is how much to increase a_ptr + # by to get the element one row down (A has M rows) + stride_am, stride_ak, + stride_bk, stride_bn, + stride_cm, stride_cn, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, + ACTIVATION: tl.constexpr, +): + """Kernel for computing the matmul C = A x B. + A has shape (M, K), B has shape (K, N) and C has shape (M, N) + """ + # ----------------------------------------------------------- + # Map program ids `pid` to the block of C it should compute. + # This is done in a grouped ordering to promote L2 data reuse + # See above `L2 Cache Optimizations` section for details + pid = tl.program_id(axis=0) + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + (pid % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + # ---------------------------------------------------------- + # Create pointers for the first blocks of A and B. + # We will advance this pointer as we move in the K direction + # and accumulate + # a_ptrs is a block of [BLOCK_SIZE_M, BLOCK_SIZE_K] pointers + # b_ptrs is a block of [BLOCK_SIZE_K, BLOCK_SIZE_n] pointers + # see above `Pointer Arithmetics` section for details + offs_am = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_bn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + offs_k = tl.arange(0, BLOCK_SIZE_K) + a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) + b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) + + # ----------------------------------------------------------- + # Iterate to compute a block of the C matrix + # We accumulate into a `[BLOCK_SIZE_M, BLOCK_SIZE_N]` block + # of fp32 values for higher accuracy. + # `accumulator` will be converted back to fp16 after the loop + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for k in range(0, K, BLOCK_SIZE_K): + # Note that for simplicity, we don't apply a mask here. + # This means that if K is not a multiple of BLOCK_SIZE_K, + # this will access out-of-bounds memory and produce an + # error or (worse!) incorrect results. + a = tl.load(a_ptrs) + b = tl.load(b_ptrs) + # We accumulate along the K dimension + accumulator += tl.dot(a, b) + # Advance the ptrs to the next K block + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K * stride_bk + # you can fuse arbitrary activation functions here + # while the accumulator is still in FP32! + if ACTIVATION == "leaky_relu": + accumulator = leaky_relu(accumulator) + c = accumulator.to(tl.float16) + + # ----------------------------------------------------------- + # Write back the block of the output matrix C + offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] + c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) + tl.store(c_ptrs, c, mask=c_mask) + + +# we can fuse `leaky_relu` by providing it as an `ACTIVATION` meta-parameter in `_matmul` +@triton.jit +def leaky_relu(x): + x = x + 1 + return tl.where(x >= 0, x, 0.01 * x) + +def matmul(a, b, activation=""): + # checks constraints + assert a.shape[1] == b.shape[0], "incompatible dimensions" + assert a.is_contiguous(), "matrix A must be contiguous" + assert b.is_contiguous(), "matrix B must be contiguous" + M, K = a.shape + K, N = b.shape + assert ( + K % 32 == 0 + ), "We don't check memory-out-of-bounds with K so K must be divisible by BLOCK_SIZE_K" + # allocates output + c = torch.empty((M, N), device=a.device, dtype=a.dtype) + # 1D launch kernel where each block gets its own program. + grid = lambda META: ( + triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv(N, META['BLOCK_SIZE_N']), + ) + matmul_kernel[grid]( + a, b, c, + M, N, K, + a.stride(0), a.stride(1), + b.stride(0), b.stride(1), + c.stride(0), c.stride(1), + ACTIVATION=activation, + ) + return c + + + +torch.manual_seed(0) +a = torch.randn((512, 512), device='cuda', dtype=torch.float16) +b = torch.randn((512, 512), device='cuda', dtype=torch.float16) +triton_output = matmul(a, b) +torch_output = torch.matmul(a, b) +print(f"triton_output={triton_output}") +print(f"torch_output={torch_output}") +if triton.testing.allclose(triton_output, torch_output): + print("✅ Triton and Torch match") +else: + print("❌ Triton and Torch differ")