diff --git a/examples/cookbook/use_layers.ipynb b/examples/cookbook/use_layers.ipynb new file mode 100644 index 00000000..ff353a53 --- /dev/null +++ b/examples/cookbook/use_layers.ipynb @@ -0,0 +1,704 @@ +{ + "nbformat": 4, + "nbformat_minor": 0, + "metadata": { + "colab": { + "provenance": [], + "gpuType": "T4", + "toc_visible": true, + "authorship_tag": "ABX9TyNmkB5Sx3PN1y+UVcpXmx/I", + "include_colab_link": true + }, + "kernelspec": { + "name": "python3", + "display_name": "Python 3" + }, + "language_info": { + "name": "python" + }, + "accelerator": "GPU" + }, + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "id": "view-in-github", + "colab_type": "text" + }, + "source": [ + "\"Open" + ] + }, + { + "cell_type": "markdown", + "source": [ + "# Kernelize your model with ease\n", + "\n", + "| Authors |\n", + "| :-- |\n", + "| [Aritra Roy Gosthipaty](https://huggingface.co/ariG23498) |\n", + "| [Sayak Paul](https://huggingface.co/sayakpaul) |\n", + "| [Daniël de Kok](https://huggingface.co/danieldk) | \n", + "\n", + "A compute kernel is a function than can be run on an accelerator (GPUs or TPUs). A kernel can be launched on many compute units (Streaming Multiprocessors in case of NVIDIA GPUs) in parallel which is how it achieves arithmetic density.\n", + "\n", + "There are two main branches while working with kernels:\n", + "1. Building kernels\n", + "2. Using built kernels\n", + "\n", + "While building kernels is an art form, in this tutorial we will showcase how we can leverage prebuilt kernels in our models. We use [`kernels` library](https://huggingface.co/docs/kernels/index), inspect kernels (from Hugging Face Hub), and use them inside a PyTorch model.\n", + "\n", + "> 🙏 This is an interactive notebook, so feel free to make changes and experiment quickly." + ], + "metadata": { + "id": "GrXqOfPzWn42" + } + }, + { + "cell_type": "markdown", + "source": [ + "## Installs and Imports" + ], + "metadata": { + "id": "Tym6Mu4PXD4N" + } + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "4rXtDdnyMmqS" + }, + "outputs": [], + "source": [ + "!pip install kernels -qq" + ] + }, + { + "cell_type": "code", + "source": [ + "import sys\n", + "import inspect\n", + "import logging\n", + "logging.basicConfig(level=logging.INFO, force=True) # important to see the kernels in action\n", + "\n", + "import torch\n", + "from torch import nn\n", + "from torch.nn import functional as F\n", + "\n", + "from kernels import get_kernel\n", + "from kernels import CUDAProperties\n", + "from kernels import LayerRepository, Device\n", + "from kernels import kernelize, Mode\n", + "from kernels import use_kernel_mapping\n", + "from kernels import use_kernelized_func\n", + "from kernels import use_kernel_forward_from_hub\n", + "from kernels import replace_kernel_forward_from_hub\n", + "\n", + "import triton\n", + "import triton.language as tl" + ], + "metadata": { + "id": "vro8PMrxP-Xs" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": [ + "## Download a kernel\n", + "\n", + "If you are looking for a model or a dataset, you go to [Hugging Face Hub](https://huggingface.co/) and select from the wide variety of options. What do you do when you are in search for a compute kernel?\n", + "\n", + "With the introduction of [`kernels`](https://huggingface.co/) you can go to the Hub, click on the [kernels tab](https://huggingface.co/kernels), and choose from the wide variety of options 🤗. Compute kernels are now a seperate repository type on the Hub.\n", + "\n", + "This means, we can:\n", + "1. Go the Hub\n", + "2. Select the kernel repository\n", + "3. Download the kernel and use it\n", + "\n", + "Let's get started!" + ], + "metadata": { + "id": "gcMFPI1ZXU1b" + } + }, + { + "cell_type": "code", + "source": [ + "activation_kernel_id = \"kernels-community/activation\"\n", + "activation_kernel_module = get_kernel(activation_kernel_id, version=1)" + ], + "metadata": { + "id": "6JJDe2_AQBBm" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": [ + "### Inspect the kernel module\n", + "\n", + "A kernel module usually holds kernel functions and kernel layers. Functions (as one would know) are input and output mapping. The idea behind a kernel function is that they are optimized for the hardware that you are using. A kernel layer on the other hand can be thought of as an optimized PyTorch module.\n", + "\n", + "Let's first inspect what the `kernels-community/activation` module holds for us." + ], + "metadata": { + "id": "XqGn9guiXYPO" + } + }, + { + "cell_type": "code", + "source": [ + "functions_available = inspect.getmembers(activation_kernel_module, inspect.isfunction)\n", + "layers_available = inspect.getmembers(activation_kernel_module.layers, inspect.isclass)\n", + "\n", + "print(\"Functions:\")\n", + "for function in functions_available:\n", + " print(f\" {function[0]}\")\n", + "\n", + "print(\"\\nLayers:\")\n", + "for layer in layers_available:\n", + " print(f\" {layer[0]}\")" + ], + "metadata": { + "id": "IzE8nHDuQK-B" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": [ + "## How to use a kernel function?\n", + "\n", + "While this cookbook (or tutorial, what ever you want to call it) is intended on showing how one would kernelize their model using a kernel layer, it is quite important to understand how the kernel functions work.\n", + "\n", + "The idea is to create the input, get the kernel function ready to be used, and then create an empty tensor that will hold the output." + ], + "metadata": { + "id": "FmQLvF97vNGP" + } + }, + { + "cell_type": "code", + "source": [ + "input = torch.randn(100, 100, device=\"cuda\", dtype=torch.float16)\n", + "kernel_output = torch.empty_like(input)\n", + "\n", + "torch_output = F.gelu(input)\n", + "activation_kernel_module.gelu(kernel_output, input)\n", + "\n", + "print(f\"Torch matches Kernel Functions: {torch.allclose(torch_output, kernel_output)}\")" + ], + "metadata": { + "id": "8sxyw6vOfywT" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": [ + "## Why do we need Kernel Layers?\n", + "\n", + "With kernel functions one would need to:\n", + "1. highlight the layer\n", + "2. replace the layer's `forward` with the kernel function\n", + "\n", + "We wanted to make this more elegant. With a kernel layer, you can just follow:\n", + "1. annotate\n", + "2. register\n", + "3. `kernelize()`\n", + "\n", + "to achieve the `forward` swap withouth having to hardcode anything. This enhances readability of code, make the intent more visible and also introduces composability.\n", + "\n", + "To see this in action, let's first build a simple model and iterate on it." + ], + "metadata": { + "id": "rT7SMHJ6XaQ9" + } + }, + { + "cell_type": "code", + "source": [ + "class MLP(nn.Module):\n", + " def __init__(self, d_in=1024, d_hidden=4096, d_out=1024):\n", + " super().__init__()\n", + " self.fc1 = nn.Linear(d_in, d_hidden)\n", + " self.act = nn.GELU()\n", + " self.fc2 = nn.Linear(d_hidden, d_out)\n", + "\n", + " def forward(self, x):\n", + " return self.fc2(self.act(self.fc1(x)))\n", + "\n", + "model = MLP().to(\"cuda\").eval()\n", + "x = torch.randn(64, 1024, device=\"cuda\")\n", + "\n", + "with torch.no_grad():\n", + " ref_out = model(x)\n", + "\n", + "print(ref_out.shape)" + ], + "metadata": { + "id": "hy-N6-MGQ2tP" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": [ + "### Build a triton kernel manually and patch it\n", + "\n", + "Here we build a very simple `gelu` kernel using [`triton`](https://triton-lang.org/main/index.html). Readers don't need to follow this code with great detail, as this is just for demonstration. It is not tuned for the GPUs being used." + ], + "metadata": { + "id": "nFSZbiCJzR68" + } + }, + { + "cell_type": "code", + "source": [ + "@triton.jit\n", + "def gelu_kernel(x_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr):\n", + " pid = tl.program_id(axis=0)\n", + " offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)\n", + " mask = offsets < n_elements\n", + "\n", + " x = tl.load(x_ptr + offsets, mask=mask)\n", + " xf = x.to(tl.float32)\n", + " y = 0.5 * xf * (1.0 + tl.math.erf(xf * 0.7071067811865476))\n", + " tl.store(out_ptr + offsets, y.to(x.dtype), mask=mask)\n", + "\n", + "\n", + "def triton_gelu(x: torch.Tensor) -> torch.Tensor:\n", + " x = x.contiguous()\n", + " out = torch.empty_like(x)\n", + " n = x.numel()\n", + " BLOCK_SIZE = 1024\n", + " grid = (triton.cdiv(n, BLOCK_SIZE),)\n", + " gelu_kernel[grid](x, out, n, BLOCK_SIZE=BLOCK_SIZE)\n", + " return out\n", + "\n", + "# sanity check against PyTorch\n", + "t = torch.randn(10_000, device=\"cuda\")\n", + "print(\"Triton matches F.gelu:\", torch.allclose(triton_gelu(t), F.gelu(t), atol=1e-5))" + ], + "metadata": { + "id": "K-FjH2Z0kSvg" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": [ + "To use this triton kernel inside our model, we would have to patch the `forward` of the `self.act` layer with the triton kernel. This is what we do next." + ], + "metadata": { + "id": "MlIUorKvzr19" + } + }, + { + "cell_type": "code", + "source": [ + "temp_model = MLP().to(\"cuda\").eval()\n", + "temp_model.load_state_dict(model.state_dict())\n", + "\n", + "temp_model.act.forward = triton_gelu # patch activation forward with the kernel\n", + "\n", + "with torch.no_grad():\n", + " triton_out = temp_model(x)\n", + "print(\"Model output matches original:\", torch.allclose(triton_out, ref_out, atol=1e-5))" + ], + "metadata": { + "id": "a_KZLR4lkkt3" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": [ + "Now let's see why this is a really clumsy way of doing things:\n", + "1. the lack of composability, everything here is hard-coded and has no leeway to compose models to our liking\n", + "2. manually monkey pathcing always leads to convoluted code bases\n", + "3. no tooling present around patching methods with kernels\n", + "\n", + "To mitigate these issues and more, we introduce kernel layers." + ], + "metadata": { + "id": "EW6tCGC8z7t_" + } + }, + { + "cell_type": "markdown", + "source": [ + "### Build a model with kernel layer\n", + "\n", + "Here we take you through the process of using kernel layers inside a model. This is a really helpful way to use kernels if you had the freedom to build the model from scratch." + ], + "metadata": { + "id": "vwbP9fas7KiK" + } + }, + { + "cell_type": "code", + "source": [ + "# 1. Annotate the layer\n", + "@use_kernel_forward_from_hub(\"Gelu\")\n", + "class Gelu(nn.Module):\n", + " def forward(self, input: torch.Tensor) -> torch.Tensor:\n", + " return F.gelu(input)\n", + "\n", + "# 2. Build the registry mapping\n", + "kernel_layer_mapping = {\n", + " \"Gelu\": {\n", + " \"cuda\": LayerRepository(\n", + " repo_id=\"kernels-community/activation\",\n", + " layer_name=\"Gelu\",\n", + " version=1,\n", + " ),\n", + " }\n", + "}\n", + "\n", + "# 3. Build the model (from source)\n", + "class MLP(nn.Module):\n", + " def __init__(self, d_in=1024, d_hidden=4096, d_out=1024):\n", + " super().__init__()\n", + " self.fc1 = nn.Linear(d_in, d_hidden)\n", + " self.act = Gelu()\n", + " self.fc2 = nn.Linear(d_hidden, d_out)\n", + "\n", + " def forward(self, x):\n", + " return self.fc2(self.act(self.fc1(x)))\n", + "\n", + "# 4. Kernelize the model\n", + "with use_kernel_mapping(kernel_layer_mapping):\n", + " kernelized_model = kernelize(MLP().to(\"cuda\"), device=\"cuda\", mode=Mode.INFERENCE)\n", + " kernelized_model.load_state_dict(model.state_dict()) # load the previous weights for equivalence testing\n", + " with torch.no_grad():\n", + " hub_out = kernelized_model(x)\n", + "\n", + "print(\"Hub kernel output matches original:\", torch.allclose(hub_out, ref_out, atol=1e-5))" + ], + "metadata": { + "id": "pQ6Qu0cC2-eo" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": [ + "We went through a bunch of steps all of which achieved the patching (as shown in the triton example above) in a more elegant way.\n", + "\n", + "1. We annotate the layer that we want to replace with a kernel layer form the Hub.\n", + "2. We then create a mapping for the layer. Here we map the layer to the kernel id and the actual layer we want to use.\n", + "3. After that we build the model using our layer.\n", + "4. Finally we `kernelize()` the model: This takes care of traversing through the model and replacing the layer's `forward` with our kernel layer from the Hub.\n", + "\n", + "If you remember the annotate, map, and kernelize formula you are good to go!\n", + "\n", + "This is quite intuitive in our opinon (but we are biased)." + ], + "metadata": { + "id": "Ku9J_VxC4zCz" + } + }, + { + "cell_type": "markdown", + "source": [ + "### Replace a layer definition\n", + "\n", + "There is another way to kernelize a model using a kernel layer. If you know the definition of the layer you can use our handy `replace_kernel_forwared_from_hub` API. This annotates the layer definition and `kernelize` will later replace it with the kernel layer." + ], + "metadata": { + "id": "YiYb-WcrqpLW" + } + }, + { + "cell_type": "code", + "source": [ + "# the model in question\n", + "print(model)" + ], + "metadata": { + "id": "7HwZzkE8csta" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# replace nn.GELU with the kernel Gelu layer\n", + "replace_kernel_forward_from_hub(nn.GELU, \"Gelu\")\n", + "\n", + "# register the mapping\n", + "kernel_layer_mapping = {\n", + " \"Gelu\": {\n", + " \"cuda\": LayerRepository(\n", + " repo_id=\"kernels-community/activation\",\n", + " layer_name=\"Gelu\",\n", + " version=1,\n", + " ),\n", + " }\n", + "}\n", + "\n", + "# kernelize the model\n", + "with use_kernel_mapping(kernel_layer_mapping):\n", + " kernelized_model = kernelize(model.to(\"cuda\"), device=\"cuda\", mode=Mode.INFERENCE)\n", + " kernelized_model.load_state_dict(model.state_dict()) # load the previous weights for equivalence testing\n", + " with torch.no_grad():\n", + " hub_out = kernelized_model(x)\n", + "\n", + "print(\"Hub kernel output matches original:\", torch.allclose(hub_out, ref_out, atol=1e-5))" + ], + "metadata": { + "id": "fqiEuMOnd9fm" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": [ + "Here we replace the `nn.GELU` torch module with the `Gelu` kernel layer we have. To do that we use the `replace_kernel_forward_from_hub` method. It is also worth noting that the method is equivalent to annotation, modifying the class.\n", + "\n", + "We strongly recommended the model developers to add `use_kernel_forward_from_hub` annotations and that `replace_kernel_forward_from_hub` is really a last resort. It relies on the upstream layer not changing in an incompatible way, whereas using the decorator implies a contract." + ], + "metadata": { + "id": "1SMngKdWegqr" + } + }, + { + "cell_type": "markdown", + "source": [ + "### Using kernel layers with functions\n", + "\n", + "There are times that you have a function that behaves exactly like the layer's `forward`. We provide a handy decorator `use_kernel_forward_from_hub` which can be used to decorate the function in concern.\n", + "\n", + "Something to note here is that decoration of functions are really for special cases. It comes in handy when you functions to kernelize, but they cannot replace the function by a layer for e.g. compatibility reasons, API breakage, etc." + ], + "metadata": { + "id": "Nf2a7-WS-IPu" + } + }, + { + "cell_type": "code", + "source": [ + "@use_kernel_forward_from_hub(\"Gelu\")\n", + "def gelu(input):\n", + " return F.gelu(input)\n", + "\n", + "@use_kernelized_func(Gelu)\n", + "class MLP(nn.Module):\n", + " def __init__(self, d_in=1024, d_hidden=4096, d_out=1024):\n", + " super().__init__()\n", + " self.fc1 = nn.Linear(d_in, d_hidden)\n", + " self.fc2 = nn.Linear(d_hidden, d_out)\n", + "\n", + " def forward(self, x):\n", + " return self.fc2(gelu(self.fc1(x))) # notice the use of the function\n", + "\n", + "\n", + "kernel_layer_mapping = {\n", + " \"Gelu\": {\n", + " \"cuda\": LayerRepository(\n", + " repo_id=\"kernels-community/activation\",\n", + " layer_name=\"Gelu\",\n", + " version=1,\n", + " ),\n", + " }\n", + "}\n", + "\n", + "with use_kernel_mapping(kernel_layer_mapping):\n", + " kernelized_model = kernelize(MLP().to(\"cuda\"), device=\"cuda\", mode=Mode.INFERENCE)\n", + " kernelized_model.load_state_dict(model.state_dict()) # load the previous weights for equivalence testing\n", + " with torch.no_grad():\n", + " hub_out = kernelized_model(x)\n", + "\n", + "print(\"Hub kernel output matches original:\", torch.allclose(hub_out, ref_out, atol=1e-5))" + ], + "metadata": { + "id": "w5Ga3bEKBuGJ" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": [ + "### Kernel Mapping\n", + "\n", + "A kernel builder wants maximum usage of their kernels. They take time developing the kernel for one hardware, and then devotes time for other hardwares too. Who are we to stop them?\n", + "\n", + "A user can map different kernels with different device specifications all at the same time. This means if you are using a model on different hardware, all the kernels that are written for the hardware follow you without having to write a single line of extra code!" + ], + "metadata": { + "id": "GgJBz-8KEK0i" + } + }, + { + "cell_type": "code", + "source": [ + "kernel_layer_mapping = {\n", + " \"Gelu\": {\n", + " \"cuda\": LayerRepository(\n", + " repo_id=\"kernels-community/activation\",\n", + " layer_name=\"Gelu\",\n", + " version=1,\n", + " ),\n", + " \"rocm\": LayerRepository(\n", + " repo_id=\"kernels-community/activation\",\n", + " layer_name=\"Gelu\",\n", + " version=1,\n", + " )\n", + " }\n", + "}" + ], + "metadata": { + "id": "L2UhLrT0dsMH" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": [ + "There are kernels for specific to inference, training, and some are also torch compilable. These are called modes. Each mode can have their own kernels. Mode is handled inside `kernelize` but we can also specify it with our kernels mapping.\n", + "\n", + "This becomes more explicit and you do not need to kernelize models differently based on the mode (as you would above). The mapping takes care of which kernel to draw from the Hub and replace with." + ], + "metadata": { + "id": "6fA3t_hWcF6C" + } + }, + { + "cell_type": "code", + "source": [ + "kernel_layer_mapping = {\n", + " \"Gelu\": {\n", + " \"cuda\": {\n", + " Mode.INFERENCE: LayerRepository(\n", + " repo_id=\"kernels-community/activation-inference-optimized\",\n", + " layer_name=\"Gelu\",\n", + " version=1,\n", + " ),\n", + " Mode.TRAINING | Mode.TORCH_COMPILE: LayerRepository(\n", + " repo_id=\"kernels-community/activation-training-optimized\",\n", + " layer_name=\"Gelu\",\n", + " version=1,\n", + " ),\n", + " }\n", + " }\n", + "}" + ], + "metadata": { + "id": "l8yUdORlcFx-" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": [ + "Another really loved feature is the CUDA Capabilites check. One can specify the CUDA compatibility and map them to the kernels that would work! These kinds of tooling makes like so much easier for the builder and the user.\n", + "\n", + "Everything is written out for you, one would just need to check the mapping and they are good to go." + ], + "metadata": { + "id": "nUDCcnrYFI1D" + } + }, + { + "cell_type": "code", + "source": [ + "kernel_layer_mapping = {\n", + " \"Gelu\": {\n", + " Device(\n", + " type=\"cuda\",\n", + " properties=CUDAProperties(\n", + " min_capability=75, max_capability=89\n", + " ),\n", + " ): LayerRepository(\n", + " repo_id=\"kernels-community/activation\",\n", + " layer_name=\"Gelu\",\n", + " version=1,\n", + " ),\n", + " Device(\n", + " type=\"cuda\",\n", + " properties=CUDAProperties(\n", + " min_capability=90, max_capability=sys.maxsize\n", + " ),\n", + " ): LayerRepository(\n", + " repo_id=\"kernels-community/activation-hopper\",\n", + " layer_name=\"Gelu\",\n", + " version=1,\n", + " ),\n", + " }\n", + "}" + ], + "metadata": { + "id": "ehbfJMVsFQ9M" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": [ + "A side note, which we thought was really interesting. If you provide an empty mapping, it de-kernelizes your model!\n", + "\n", + "```\n", + "with use_kernel_mapping({}):\n", + " undo_kernelized_model = kernelize(kernelized_model, device=\"cuda\", mode=Mode.INFERENCE)\n", + " with torch.no_grad():\n", + " undo_hub_out = undo_kernelized_model(x)\n", + "```" + ], + "metadata": { + "id": "mK0g1l9FUqF8" + } + }, + { + "cell_type": "markdown", + "source": [ + "## Are kernels useful?\n", + "\n", + "We know this should be at the very top of the cookbook. But knowing that you have read it through, we assume that you already know what kernels get you. In this section we just showcase a simple case where we generate some tokens with a model with and without kernels.\n", + "\n", + "Here we showcase the time required to generate 256 tokens for `gpt-oss-2b`. We benchmark the model with and without using kernels layers. We also showcase how more batch size affects the benchmark.\n", + "\n", + "| ![gpt-oss](https://huggingface.co/datasets/ariG23498/colab-images/resolve/main/gpt-kernels-plot.png) |\n", + "| :--: |\n", + "| Benchmark `gpt-oss-20b` with and without kernels |\n", + "\n", + "You can reproduce the results:\n", + "1. [Benchmarking script](https://gist.github.com/ariG23498/b2f6c78ee40d5a02119566d58d00fc56)\n", + "2. Run it with Hugging Face Jobs: `uvx hf jobs uv run --flavor h200 benchmark_gpt.py -d`\n", + "3. See the [Jobs live](https://huggingface.co/jobs/ariG23498/6ac7c9f5fee2c9007016b027)" + ], + "metadata": { + "id": "m3zgDv9jI5zg" + } + }, + { + "cell_type": "markdown", + "source": [ + "## Conclusion\n", + "\n", + "We hope this interactive cookbook gets you up to speed with how one would use kernels which are built and uploaded to the Hub. We hope to seeing more of you using kernels for your models and also your libraries.\n", + "\n", + "Kernels are really powerful and we want you to use them leaving all the tooling issues for us to tackle! Keep the hardware accelerators warm 🔥." + ], + "metadata": { + "id": "4rnNhS_5FlR9" + } + } + ] +} \ No newline at end of file