tutorial/notebook/Petals_with_Llama2_and_Stable_Beluga2.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "<div align=\"center\">\n",
        "<img src=\"https://camo.githubusercontent.com/473dd9f992924d27457650251786464f72e54121ac6e9210add0f483ca849277/68747470733a2f2f692e696d6775722e636f6d2f3765523750616e2e706e67\" width=\"40%\">  \n",
        "</div>\n",
        "\n",
        "# Getting started with Petals\n",
        "\n",
        "This notebook will guide you through the basics of Petals &mdash; a system for inference and fine-tuning large language models without the need to have high-end GPUs. With Petals, you can join compute resources with other people over the Internet and run large language models such as Llama 2 (70B), Stable Beluga 2, Llama-65B, Guanaco-65B, or BLOOM-176B right from your desktop computer or Google Colab.\n",
        "\n",
        "💬 If you meet any issues while running this notebook, let us know in the **[#running-a-client](https://discord.gg/J29mCBNBvm)** channel of our Discord!\n",
        "\n",
        "So, let's get started! First, let's install [the Petals package](https://github.com/bigscience-workshop/petals):"
      ],
      "metadata": {
        "id": "VsXHWJLuowcn"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pBC52TF3LVY1"
      },
      "outputs": [],
      "source": [
        "%pip install git+https://github.com/bigscience-workshop/petals"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Step 1. Loading the distributed model 🚀\n",
        "\n",
        "Let's start with the easiest task &mdash; creating a distributed model and using it for generating text. This machine will download a small part of the model weights and rely on other computers in the network to run the rest of the model.\n",
        "\n",
        "The Petals interface is compatible with PyTorch and the 🤗 [Transformers](https://github.com/huggingface/transformers) library &mdash; it feels like you're working with a local model even though its parts are hosted remotely. We suggest to start with [Stable Beluga 2 (70B)](https://huggingface.co/stabilityai/StableBeluga2), one of the best fine-tuned variants of Llama 2, but you can also use standard [Llama&nbsp;2&nbsp;(70B)](https://huggingface.co/meta-llama/Llama-2-70b-hf) if you have access to it (see below)."
      ],
      "metadata": {
        "id": "yEbot-oEXdpw"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2uuX1IMLLotQ"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "from transformers import AutoTokenizer\n",
        "from petals import AutoDistributedModelForCausalLM\n",
        "\n",
        "model_name = \"petals-team/StableBeluga2\"\n",
        "# You could also use \"meta-llama/Llama-2-70b-chat-hf\" or any other supported model from 🤗 Model Hub\n",
        "\n",
        "tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False, add_bos_token=False)\n",
        "model = AutoDistributedModelForCausalLM.from_pretrained(model_name)\n",
        "model = model.cuda()"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "🦙 **Want to run Llama 2?** Request access to its weights at the ♾️ [Meta AI website](https://ai.meta.com/resources/models-and-libraries/llama-downloads/) and 🤗 [Model Hub](https://huggingface.co/meta-llama/Llama-2-70b-hf) (make sure to use the same email),  get an 🔑 [access token](https://huggingface.co/settings/tokens), then run `!huggingface-cli login --token YOUR_TOKEN` before loading the model. Or just try it in our [chatbot app](https://chat.petals.dev).\n",
        "\n",
        "📋 **Friendly reminder.** This Colab is provided for demo purposes. If you want to use these models in your own projects, make sure you follow their terms of use (see the ones for [Stable Beluga 2](https://huggingface.co/stabilityai/StableBeluga2/blob/main/LICENSE.txt) and [Llama 2](https://bit.ly/llama2-license))."
      ],
      "metadata": {
        "id": "JhKyusDDovS1"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### ✍️ How to generate text?\n",
        "\n",
        "Let's try to generate something by calling __`model.generate()`__ method.\n",
        "\n",
        "The first call to this method takes a few seconds to connect to the Petals swarm. Once we do that, you should expect generation speed of up to **5-6 tokens/sec**. If you don't have enough GPU memory to host the entire model, this is much faster than what you get with other methods, such as offloading or running the model on CPU."
      ],
      "metadata": {
        "id": "zhyUxv13sfKO"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "inputs = tokenizer('A cat in French is \"', return_tensors=\"pt\")[\"input_ids\"].cuda()\n",
        "outputs = model.generate(inputs, max_new_tokens=3)\n",
        "print(tokenizer.decode(outputs[0]))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1s1IrE1H8wwr",
        "outputId": "1e40b217-0ae3-47e4-b893-d56ed8be1f39"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "A cat in French is \"chat\" and\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "The `model.generate()` method runs **greedy** generation by default, but you can use many other generation methods like **top-p/top-k sampling** or **beam search** &mdash; just set proper arguments for the 🤗 Transformers [.generate()](https://huggingface.co/blog/how-to-generate) method.\n",
        "\n",
        "🔏 **Note:** Your data is processed by other people in the public swarm. Learn more about privacy [here](https://github.com/bigscience-workshop/petals/wiki/Security,-privacy,-and-AI-safety). For sensitive data, you can set up a [private swarm](https://github.com/bigscience-workshop/petals/wiki/Launch-your-own-swarm) among people you trust."
      ],
      "metadata": {
        "id": "02d0BDEAuUFQ"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Step 2. Generating tokens on the fly and making chatbots 🕊️\n",
        "\n",
        "If you'd like to talk to the model in an interactive way, you can use the __inference session__ interface &mdash; it allows to print generated tokens on the fly or make a chat bot that responds to human's inputs.\n",
        "\n",
        "The inference session looks for a sequence of servers to run successive inference steps and store past attention caches. This way, you don't need to rerun previous tokens through the transformer to generate each phrase. If one of the servers disconnects or fails, Petals will automatically find a replacement and regenerate only a small part of the caches.\n",
        "\n",
        "We provide a scheme of what's going on below, and you can check out servers that are actually online right now in our 🏥 [health monitor](https://health.petals.dev):\n",
        "\n",
        "<br>\n",
        "<div align=\"center\">\n",
        "<img src=\"https://i.imgur.com/fKR9BSP.png\" width=\"70%\">\n",
        "</div>\n",
        "\n",
        "Let's see how to show tokens on the fly, as soon as they are generated:"
      ],
      "metadata": {
        "id": "fZlzYVn0ApyF"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "fake_token = tokenizer(\"^\")[\"input_ids\"][0]  # Workaround to make tokenizer.decode() keep leading spaces\n",
        "\n",
        "text = \"What is a good chatbot? Answer:\"\n",
        "prefix = tokenizer(text, return_tensors=\"pt\")[\"input_ids\"].cuda()\n",
        "\n",
        "with model.inference_session(max_length=30) as sess:\n",
        "    for i in range(20):\n",
        "        # Prefix is passed only for the 1st token of the outputs\n",
        "        inputs = prefix if i == 0 else None\n",
        "\n",
        "        # Let's use sampling with temperature = 0.9 and top_p = 0.6 to get more diverse results\n",
        "        outputs = model.generate(inputs, max_new_tokens=1, session=sess,\n",
        "                                 do_sample=True, temperature=0.9, top_p=0.6)\n",
        "\n",
        "        text += tokenizer.decode([fake_token, outputs[0, -1].item()])[1:]\n",
        "        print(text)"
      ],
      "metadata": {
        "id": "lTovfF4HxOuJ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Writing a simple chatbot 💬\n",
        "\n",
        "Now, let's proceed to writing a simple chatbot! We'll need one more loop that accepts inputs from a human, then runs generation until we get a newline (**`\\n`**):"
      ],
      "metadata": {
        "id": "lgzYdGVzyR1y"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "with model.inference_session(max_length=512) as sess:\n",
        "    while True:\n",
        "        prompt = input('Human: ')\n",
        "        if prompt == \"\":\n",
        "            break\n",
        "        prefix = f\"Human: {prompt}\\nFriendly AI:\"\n",
        "        prefix = tokenizer(prefix, return_tensors=\"pt\")[\"input_ids\"].cuda()\n",
        "        print(\"Friendly AI:\", end=\"\", flush=True)\n",
        "\n",
        "        while True:\n",
        "            outputs = model.generate(prefix, max_new_tokens=1, session=sess,\n",
        "                                     do_sample=True, temperature=0.9, top_p=0.6)\n",
        "            outputs = tokenizer.decode([fake_token, outputs[0, -1].item()])[1:]\n",
        "\n",
        "            # Now, let's print one new token at a time\n",
        "            print(outputs, end=\"\", flush=True)\n",
        "\n",
        "            if \"\\n\" in outputs or \"</s>\" in outputs:\n",
        "                break\n",
        "            prefix = None  # Prefix is passed only for the 1st token of the bot's response"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "m5bdMaYYAqYR",
        "outputId": "8ece4869-da31-45ae-b3f7-daf7d69c253c"
      },
      "execution_count": null,
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Human: Hi, how are you?\n",
            "Friendly AI: I'm fine, thanks. And you?\n",
            "Human: \n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "📦 **Deploying apps that use Petals.** You can wrap up your code into a web app available for other people. As an example, take a look at our 💬&nbsp;[chatbot web app](https://chat.petals.dev) that uses this 🗼 [HTTP endpoint](https://github.com/petals-infra/chat.petals.dev#apis) for inference under the hood."
      ],
      "metadata": {
        "id": "39Ym5TZ_ZWUd"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Step 3. How does it work? 🛠️\n",
        "\n",
        "The `model` you are running is equal to the original model, but only a part of it is loaded into your machine's GPU. Let's have a look under the hood:"
      ],
      "metadata": {
        "id": "557VBCGS8Dpv"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "model"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "FB-lCjKs8NEt",
        "outputId": "68be5b2d-05db-41d9-a9b9-d98826a13518"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DistributedLlamaForCausalLM(\n",
              "  (model): DistributedLlamaModel(\n",
              "    (embed_tokens): Embedding(32000, 8192, padding_idx=0)\n",
              "    (layers): RemoteSequential(modules=llama-65b-hf.0..llama-65b-hf.79)\n",
              "    (norm): LlamaRMSNorm()\n",
              "  )\n",
              "  (lm_head): LMHead()\n",
              ")"
            ]
          },
          "metadata": {},
          "execution_count": 18
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "As you can see, word embeddings and some other layers are regular PyTorch modules hosted on your machine, but the rest of the model (e.g., transformers blocks) is encased in the __`RemoteSequential`__ class. This is an advanced PyTorch module that runs on a distributed swarm of other machines.\n",
        "\n",
        "Still, you can access individual layers and their outputs, as well as run forward/backward through them:"
      ],
      "metadata": {
        "id": "9EW3wBDO-aTJ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "first_five_layers = model.model.layers[0:5]\n",
        "first_five_layers"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "OHYihUHH-Zo6",
        "outputId": "1acfc043-50f8-464b-d5e4-d2ee8c8151df"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "RemoteSequential(modules=llama-65b-hf.0..llama-65b-hf.4)"
            ]
          },
          "metadata": {},
          "execution_count": 19
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "dummy_inputs = torch.randn(1, 3, model.config.hidden_size, dtype=torch.bfloat16, requires_grad=True)\n",
        "outputs = first_five_layers(dummy_inputs)\n",
        "outputs"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ilOwoxr47Sso",
        "outputId": "fda0e151-b6f8-4d3c-b11f-12319b38f532"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([[[-0.2139,  0.4180, -2.1562,  ..., -0.2031, -0.6523,  1.9844],\n",
              "         [-1.8594,  0.7070,  0.9141,  ...,  1.2109,  1.4219, -0.6680],\n",
              "         [ 1.0469,  0.0845, -0.9609,  ...,  1.6406,  1.3672,  1.1875]]],\n",
              "       dtype=torch.bfloat16, grad_fn=<_RemoteSequentialAutogradFunctionBackward>)"
            ]
          },
          "metadata": {},
          "execution_count": 24
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "loss = torch.mean((outputs - torch.ones_like(outputs)) ** 2)\n",
        "loss.backward()  # backpropagate through the internet\n",
        "print(\"Grad w.r.t. inputs:\", dummy_inputs.grad.flatten())"
      ],
      "metadata": {
        "id": "4XM6In8uArE5",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "dbcbf8e1-cd51-4f18-e04e-677fd7f9004e"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Grad w.r.t. inputs: tensor([-2.2507e-04, -1.9550e-05, -4.4250e-04,  ...,  5.8174e-05,\n",
            "         5.0545e-05,  5.9605e-05], dtype=torch.bfloat16)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "You can see that PyTorch can calculate gradients through remote blocks, which we'll later use to run fine-tuning!\n",
        "\n",
        "Note that, in general, you can mix and match distributed layers like in regular PyTorch and even insert and train your own layers (e.g., adapters) between the pre-trained ones. You can find further technical details in our [research paper](https://arxiv.org/pdf/2209.01188.pdf)."
      ],
      "metadata": {
        "id": "DxWOzWCwat6H"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Step 4. Prompt-tuning, or making a fox innocent 🦊\n",
        "\n",
        "While the remotely hosted transformer blocks are frozen to keep the pretrained model the same for all users, using **parameter-efficient fine-tuning** methods like trainable prompts or adapters (e.g., [LoRA](https://arxiv.org/abs/2106.09685)) is usually enough to make the model solve most downstream tasks. This way, all trainable parameters and the optimizer will be hosted locally, so you'd be able to fine-tune the model without interfering with other users.\n",
        "\n",
        "In this section, we'll use **trainable prompts** to solve a dummy task &mdash; take the model saying \"*A quick brown fox jumps over the lazy dog.*\" and teach it to say the opposite &ndash; that actually \"*A quick brown fox did not jump over the lazy dog*\".\n",
        "\n",
        "Let's see how the off-the-shelf model behaves on this task:"
      ],
      "metadata": {
        "id": "miR8kS64ygvE"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "inputs = tokenizer(\"A quick brown fox\", return_tensors=\"pt\")[\"input_ids\"].cuda()\n",
        "outputs = model.generate(inputs, max_new_tokens=7)\n",
        "print(\"generated:\", tokenizer.decode(outputs[0]))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "EHuVt22DzVvV",
        "outputId": "8c98885f-fc35-4d5e-877a-b684a4399434"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "generated: A quick brown fox jumps over the lazy dog.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "If you aren't familiar with prompt tuning, **trainable prompts** are just a few trainable \"tokens\" added before the inputs of your model. **Deep prompt tuning** adds extra trainable \"tokens\" for each transformer block &mdash; this way we'll have more trainable parameters:\n",
        "\n",
        "<div align=\"center\">\n",
        "<img src=\"https://i.imgur.com/eohYuNE.png\" width=\"60%\">\n",
        "</div>\n",
        "\n",
        "**Figure.** The scheme of **deep prompt tuning** from Liu, Xiao, et al. [\"P-tuning v2: Prompt tuning can be comparable to fine-tuning universally across scales and tasks.\"](https://arxiv.org/abs/2110.07602).\n",
        "\n",
        "Petals supports prompt tuning (**`tuning_mode=\"ptune\"`**) and deep prompt tuning (**`tuning_mode=\"deep_ptune\"`**) out of the box. Let's use deep prompt tuning with 3 tokens for input and each transformer block (**`pre_seq_len=3`**):"
      ],
      "metadata": {
        "id": "Nlg-SkpLe9Ly"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "model = AutoDistributedModelForCausalLM.from_pretrained(model_name, tuning_mode='deep_ptune', pre_seq_len=3)\n",
        "model = model.cuda()"
      ],
      "metadata": {
        "id": "qLsi50_l9UMS",
        "outputId": "a49f36ec-6168-494e-a8a9-a2e8b34ae8f0",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Jun 23 15:12:42.335 [\u001b[1m\u001b[34mINFO\u001b[0m] Using DHT prefix: llama-65b-hf\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now, we can create an Adam optimizer and fine-tune the distributed model &mdash; all done in the same way as if you were training a local model:"
      ],
      "metadata": {
        "id": "aYaE8rxAf7o8"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "opt = torch.optim.Adam(model.parameters(), lr=1e-3)\n",
        "\n",
        "the_fox_is_innocent = tokenizer(\"A quick brown fox did not jump over the lazy dog\", return_tensors=\"pt\")[\"input_ids\"].cuda()\n",
        "for i in range(12):\n",
        "    loss = model(input_ids=the_fox_is_innocent, labels=the_fox_is_innocent).loss\n",
        "    print(f\"loss[{i}] = {loss.item():.3f}\")\n",
        "\n",
        "    opt.zero_grad()\n",
        "    loss.backward()\n",
        "    opt.step()\n",
        "    print(\"opt.step()\")"
      ],
      "metadata": {
        "id": "Zm2eWZew9b3T"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Once the loss function is close to zero, you can see that the model continues the sentence as we want it to:"
      ],
      "metadata": {
        "id": "L83Ix3mZgI3V"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "inputs = tokenizer(\"A quick brown fox\", return_tensors=\"pt\")[\"input_ids\"].cuda()\n",
        "outputs = model.generate(inputs, max_new_tokens=7)\n",
        "print(\"generated:\", tokenizer.decode(outputs[0]))"
      ],
      "metadata": {
        "id": "EZYRU7v_A5DO",
        "outputId": "24bc9dbb-7bd7-4290-e36c-6f6249e88073",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "generated: A quick brown fox did not jump over the lazy dog\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Step 5. Fine-tuning a trainable adapter 🏋️\n",
        "\n",
        "In this section, we'll try another popular fine-tuning method &mdash; adding **trainable adapters** to the model. These are small trainable layers added between the pretrained transformer blocks or in addition to some block weights.\n",
        "\n",
        "Here, we'll add a basic **trainable** linear layer in the middle of the pretrained model and swap the model's head to perform a **classification task**\n",
        "instead of generating text. As earlier, this layer's weights and the corresponding optimizer statistics will be stored locally:"
      ],
      "metadata": {
        "id": "lJt4OS2pIZe_"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "\n",
        "model = AutoDistributedModelForCausalLM.from_pretrained(model_name)\n",
        "model = model.cuda()\n",
        "\n",
        "class LLMBasedClassifier(nn.Module):\n",
        "  def __init__(self, model):\n",
        "    super().__init__()\n",
        "    self.distributed_layers = model.transformer.h\n",
        "    self.adapter = nn.Sequential(nn.Linear(model.config.hidden_size, 32), nn.Linear(32, model.config.hidden_size))\n",
        "    self.head = nn.Linear(model.config.hidden_size, 2)\n",
        "\n",
        "  def forward(self, embeddings):\n",
        "    mid_block = len(self.distributed_layers) // 2\n",
        "    hidden_states = self.distributed_layers[:mid_block](embeddings)\n",
        "    hidden_states = self.adapter(hidden_states)\n",
        "    hidden_states = self.distributed_layers[mid_block:](hidden_states)\n",
        "    pooled_states = torch.mean(hidden_states, dim=1)\n",
        "    return self.head(pooled_states)"
      ],
      "metadata": {
        "id": "YHPFitSoIZO0"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now, let's take the Adam optimizer and train the model with a cross-entropy loss, which is typically used for classification:"
      ],
      "metadata": {
        "id": "HS2xbjsChZus"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "classifier = LLMBasedClassifier(model).cuda()\n",
        "opt = torch.optim.Adam(classifier.parameters(), 3e-5)\n",
        "inputs = torch.randn(3, 2, model.config.hidden_size, device='cuda')\n",
        "labels = torch.tensor([1, 0, 1], device='cuda')\n",
        "\n",
        "for i in range(5):\n",
        "  loss = F.cross_entropy(classifier(inputs), labels)\n",
        "  print(f\"loss[{i}] = {loss.item():.3f}\")\n",
        "  opt.zero_grad()\n",
        "  loss.backward()\n",
        "  opt.step()\n",
        "\n",
        "print('predicted:', classifier(inputs).argmax(-1))  # l, o, l"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "OJ7AFOumIYlY",
        "outputId": "dffa3325-fd0c-4a5d-c9e4-d1c1ddabfd24"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "loss[0] = 11.039\n",
            "loss[1] = 6.550\n",
            "loss[2] = 2.489\n",
            "loss[3] = 0.455\n",
            "loss[4] = 0.038\n",
            "predicted: tensor([1, 0, 1], device='cuda:0')\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "You can see that the loss is decreasing &mdash; that is, the model overfits to our dummy dataset! If you're further interested in a full-fledged example of fine-tuning on a classification task, check out [this notebook](https://colab.research.google.com/github/bigscience-workshop/petals/blob/main/examples/prompt-tuning-sst2.ipynb) where we fine-tune Llama on the popular SST2 dataset."
      ],
      "metadata": {
        "id": "wi6EZGRDhxqQ"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Step 6. Using custom sampling methods 🎰\n",
        "\n",
        "In **Step 3**, you've seen that you can write normal PyTorch code to work with the parts of our distributed model. In practice, you can use it to implement many tricky fine-tuning and sampling methods &mdash; something you can't usually do with hosted APIs.\n",
        "\n",
        "Let's show how to implement your own **sampling method** from scratch. The __`model.inference_session()`__ interface in Petals allows you to write custom inference code. You can use this to implement any sampling algorithms you want, as well as write a custom beam search algorithm (e.g., to forbid the model from using swearwords).\n",
        "\n",
        "Below, we reimplement the standard `model.generate()` interface by making forward passes through all the layers manually:"
      ],
      "metadata": {
        "id": "xz6EJ8VsYW2b"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "frLiu0yiL6Sn",
        "outputId": "4064d246-f2f4-4c4b-e074-3fab62ab194c"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Jun 24 02:59:32.990 [\u001b[1m\u001b[34mINFO\u001b[0m] What is a good chatbot? Answer: A\n",
            "Jun 24 02:59:34.073 [\u001b[1m\u001b[34mINFO\u001b[0m] What is a good chatbot? Answer: A chat\n",
            "Jun 24 02:59:35.088 [\u001b[1m\u001b[34mINFO\u001b[0m] What is a good chatbot? Answer: A chatbot\n",
            "Jun 24 02:59:36.103 [\u001b[1m\u001b[34mINFO\u001b[0m] What is a good chatbot? Answer: A chatbot that\n",
            "Jun 24 02:59:37.165 [\u001b[1m\u001b[34mINFO\u001b[0m] What is a good chatbot? Answer: A chatbot that is\n",
            "Jun 24 02:59:38.200 [\u001b[1m\u001b[34mINFO\u001b[0m] What is a good chatbot? Answer: A chatbot that is able\n",
            "Jun 24 02:59:39.219 [\u001b[1m\u001b[34mINFO\u001b[0m] What is a good chatbot? Answer: A chatbot that is able to\n",
            "Jun 24 02:59:40.245 [\u001b[1m\u001b[34mINFO\u001b[0m] What is a good chatbot? Answer: A chatbot that is able to understand\n",
            "Jun 24 02:59:41.226 [\u001b[1m\u001b[34mINFO\u001b[0m] What is a good chatbot? Answer: A chatbot that is able to understand the\n",
            "Jun 24 02:59:42.229 [\u001b[1m\u001b[34mINFO\u001b[0m] What is a good chatbot? Answer: A chatbot that is able to understand the user\n",
            "Jun 24 02:59:43.208 [\u001b[1m\u001b[34mINFO\u001b[0m] What is a good chatbot? Answer: A chatbot that is able to understand the user’\n",
            "Jun 24 02:59:44.250 [\u001b[1m\u001b[34mINFO\u001b[0m] What is a good chatbot? Answer: A chatbot that is able to understand the user’s\n",
            "Jun 24 02:59:45.249 [\u001b[1m\u001b[34mINFO\u001b[0m] What is a good chatbot? Answer: A chatbot that is able to understand the user’s intent\n",
            "Jun 24 02:59:46.213 [\u001b[1m\u001b[34mINFO\u001b[0m] What is a good chatbot? Answer: A chatbot that is able to understand the user’s intent and\n",
            "Jun 24 02:59:47.206 [\u001b[1m\u001b[34mINFO\u001b[0m] What is a good chatbot? Answer: A chatbot that is able to understand the user’s intent and provide\n",
            "Jun 24 02:59:48.198 [\u001b[1m\u001b[34mINFO\u001b[0m] What is a good chatbot? Answer: A chatbot that is able to understand the user’s intent and provide the\n"
          ]
        }
      ],
      "source": [
        "from hivemind import get_logger\n",
        "\n",
        "logger = get_logger()\n",
        "\n",
        "fake_token = tokenizer(\"^\")[\"input_ids\"][0]  # Workaround to make tokenizer.decode() keep leading spaces\n",
        "\n",
        "text = \"What is a good chatbot? Answer:\"\n",
        "token_ids = tokenizer(text, return_tensors=\"pt\")[\"input_ids\"].cuda()\n",
        "max_length = 100\n",
        "with torch.inference_mode():\n",
        "    with model.inference_session(max_length=max_length) as sess:\n",
        "        while len(text) < max_length:\n",
        "            embs = model.transformer.word_embeddings(token_ids)\n",
        "            embs = model.transformer.word_embeddings_layernorm(embs)\n",
        "\n",
        "            h = sess.step(embs)\n",
        "            h_last = model.transformer.ln_f(h[:, -1])\n",
        "            logits = model.lm_head(h_last)\n",
        "\n",
        "            next_token = logits.argmax(dim=-1)\n",
        "            text += tokenizer.decode([fake_token, next_token.item()])[1:]\n",
        "            token_ids = next_token.reshape(1, 1)\n",
        "            logger.info(text)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Step 7. Sharing is caring 🤗\n",
        "\n",
        "Petals is a community-run system &mdash; we rely on people sharing their GPUs. You can check out available servers on our [swarm monitor](https://health.petals.dev) and connect your GPU to help serving one of the models!\n",
        "\n",
        "🐍 **Linux + Anaconda.** Run these commands:\n",
        "\n",
        "```bash\n",
        "conda install pytorch pytorch-cuda=11.7 -c pytorch -c nvidia\n",
        "pip install git+https://github.com/bigscience-workshop/petals\n",
        "python -m petals.cli.run_server petals-team/StableBeluga2\n",
        "```\n",
        "\n",
        "🪟 **Windows + WSL.** Follow the guide on our [Wiki](https://github.com/bigscience-workshop/petals/wiki/Run-Petals-server-on-Windows).\n",
        "\n",
        "🐋 **Any OS + Docker.** Run our [Docker](https://www.docker.com) image:\n",
        "\n",
        "```bash\n",
        "sudo docker run -p 31330:31330 --ipc host --gpus all --volume petals-cache:/cache --rm learningathome/petals:main \\\n",
        "    python -m petals.cli.run_server --port 31330 petals-team/StableBeluga2\n",
        "```\n",
        "\n",
        "These commands will host a part of [Stable Beluga 2](https://huggingface.co/stabilityai/StableBeluga2) on your machine. You can also host `meta-llama/Llama-2-70b-hf`, `meta-llama/Llama-2-70b-chat-hf`, repos with Llama-65B, `bigscience/bloom`, `bigscience/bloomz`, and other compatible models from 🤗 [Model Hub](https://huggingface.co/models), or [add support](https://github.com/bigscience-workshop/petals/wiki/Run-a-custom-model-with-Petals) for new model architectures.\n",
        "\n",
        "🦙 **Want to host Llama 2?** Request access to its weights at the ♾️ [Meta AI website](https://ai.meta.com/resources/models-and-libraries/llama-downloads/) and 🤗 [Model Hub](https://huggingface.co/meta-llama/Llama-2-70b-hf), generate an 🔑 [access token](https://huggingface.co/settings/tokens), then use this command for `petals.cli.run_server`:\n",
        "\n",
        "```bash\n",
        "python -m petals.cli.run_server meta-llama/Llama-2-70b-chat-hf --token YOUR_TOKEN_HERE\n",
        "```\n",
        "\n",
        "💬 **FAQ.** Check out our [Wiki](https://github.com/bigscience-workshop/petals/wiki/FAQ:-Frequently-asked-questions#running-a-server) to learn how to use multple GPUs, restart the server on reboot, etc. If you have any issues, ping us in [our Discord](https://discord.gg/D9MwApKgWa)!\n",
        "\n",
        "🔒 **Security.** Hosting a server does not allow others to run custom code on your computer. Learn more [here](https://github.com/bigscience-workshop/petals/wiki/Security,-privacy,-and-AI-safety).\n",
        "\n",
        "🏆 **Thank you!** Once you load and host 10+ blocks, we can show your name or link on the [swarm monitor](https://health.petals.dev) as a way to say thanks. You can specify them with `--public_name YOUR_NAME`."
      ],
      "metadata": {
        "id": "iK_iT8J3zDC0"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## What's next?\n",
        "\n",
        "Congratulations on finishing our tutorial! You've learned how to use Petals for different tasks, how it works under the hood, and how to increase its capacity.\n",
        "\n",
        "You can find a few other helpful resources below:\n",
        "\n",
        "* __More about Petals.__ The [README](https://github.com/bigscience-workshop/petals#readme) file in our GitHub repository has links to more Petals-related materials, including instructions for starting a private swarm and adding new models.\n",
        "\n",
        "* __Discord server.__ If you have any feedback, questions, or technical issues, please [join our Discord server](https://discord.gg/D9MwApKgWa) and let us know. If you want to build something based on Petals, we'd be happy to hear what you are up to.\n",
        "\n",
        "* __Research paper.__ Our [paper](https://arxiv.org/abs/2209.01188) shares more details about our research and what happens in Petals under the hood."
      ],
      "metadata": {
        "id": "PGUpxhQxXVCF"
      }
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": []
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}