examples/modify_existing.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we demonstrate the workflow of adding sparsity into existing models. As an example we take a\n",
    "single encoder layer of BERT."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using cache found in /users/aivanov/.cache/torch/hub/huggingface_pytorch-transformers_main\n",
      "Some weights of the model checkpoint at bert-base-uncased were not used when initializing BertModel: ['cls.seq_relationship.bias', 'cls.seq_relationship.weight', 'cls.predictions.transform.LayerNorm.bias', 'cls.predictions.transform.dense.weight', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.transform.dense.bias', 'cls.predictions.decoder.weight', 'cls.predictions.bias']\n",
      "- This IS expected if you are initializing BertModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing BertModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([8, 128, 768])\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "input_shape = (8, 128, 768) # batch, sequence, features\n",
    "model = torch.hub.load('huggingface/pytorch-transformers',\n",
    "    'model', 'bert-base-uncased').encoder.layer[0]\n",
    "input = torch.rand(input_shape)\n",
    "output = model(input)\n",
    "print(output[0].shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We target all linear layers in this model, including feedforward and attention projection layers.\n",
    "A linear layer computes $y = xA^T + b$ and is defined in the `torch.nn.Linear` module.\n",
    "In particular, we are going to sparsify tensors $A$ by magnitude pruning of $90\\%$ of their values and storing them in the CSR format.\n",
    "In the following snippet we collect the six weight tensors from linear layers, and assign sparsifiers to them. This yields the fully qualified names assigned by PyTorch to each of these tensors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/users/aivanov/miniconda3/envs/sten1/lib/python3.9/site-packages/torch/distributed/distributed_c10d.py:181: UserWarning: torch.distributed.reduce_op is deprecated, please use torch.distributed.ReduceOp instead\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['attention.self.query.weight', 'attention.self.key.weight', 'attention.self.value.weight', 'attention.output.dense.weight', 'intermediate.dense.weight', 'output.dense.weight']\n"
     ]
    }
   ],
   "source": [
    "import sten\n",
    "weights_to_sparsify = []\n",
    "sb = sten.SparsityBuilder()\n",
    "for module_name, module in model.named_modules():\n",
    "    if isinstance(module, torch.nn.modules.linear.Linear):\n",
    "        weight = module_name + \".weight\"\n",
    "        weights_to_sparsify.append(weight)\n",
    "        sb.set_weight(\n",
    "            name=weight,\n",
    "            initial_sparsifier=sten.ScalarFractionSparsifier(0.9),\n",
    "            inline_sparsifier=sten.KeepAll(),\n",
    "            tmp_format=torch.Tensor,\n",
    "            external_sparsifier=sten.KeepAll(),\n",
    "            out_format=sten.CsrTensor,\n",
    "        )\n",
    "print(weights_to_sparsify)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we repeat the same process for intermediate tensors.\n",
    "In this example, we target only the output of the GELU activation.\n",
    "However, it is challenging to refer to this intermediate tensor, as we treat the module as a black box that we do not modify, and internal operators may have varying or no name, depending on the implementation.\n",
    "Examining the layer modules shows the model structure:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "BertLayer(\n",
      "  (attention): BertAttention(\n",
      "    (self): BertSelfAttention(\n",
      "      (query): Linear(in_features=768, out_features=768, bias=True)\n",
      "      (key): Linear(in_features=768, out_features=768, bias=True)\n",
      "      (value): Linear(in_features=768, out_features=768, bias=True)\n",
      "      (dropout): Dropout(p=0.1, inplace=False)\n",
      "    )\n",
      "    (output): BertSelfOutput(\n",
      "      (dense): Linear(in_features=768, out_features=768, bias=True)\n",
      "      (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
      "      (dropout): Dropout(p=0.1, inplace=False)\n",
      "    )\n",
      "  )\n",
      "  (intermediate): BertIntermediate(\n",
      "    (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
      "    (intermediate_act_fn): GELUActivation()\n",
      "  )\n",
      "  (output): BertOutput(\n",
      "    (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
      "    (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
      "    (dropout): Dropout(p=0.1, inplace=False)\n",
      "  )\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "From this we see that the `model.intermediate` submodule contains the GELU activation, but we still do not know the name of the output intermediate tensor.\n",
    "We use the `torch.fx` tracer to assign deterministic names to the intermediate tensors. The result of running this command shows that the output of `<built-in function gelu>` (accessible as `torch.nn.functional.gelu`) is assigned to the tensor with the name `gelu` inside the `model.intermediate` module."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "opcode         name           target                    args              kwargs\n",
      "-------------  -------------  ------------------------  ----------------  --------\n",
      "placeholder    hidden_states  hidden_states             ()                {}\n",
      "call_module    dense          dense                     (hidden_states,)  {}\n",
      "call_function  gelu           <built-in function gelu>  (dense,)          {}\n",
      "output         output         output                    (gelu,)           {}\n"
     ]
    }
   ],
   "source": [
    "torch.fx.symbolic_trace(model.intermediate).graph.print_tabular()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now assign a random fraction sparsifier with $90\\%$ zeroing probability to the GELU output intermediate tensor.\n",
    "The sparsifier stores the tensor in COO format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "sb.set_interm(\n",
    "    name=\"intermediate.gelu\",\n",
    "    inline_sparsifier=sten.RandomFractionSparsifier(0.9),\n",
    "    tmp_format=sten.CooTensor,\n",
    "    external_sparsifier=sten.KeepAll(),\n",
    "    out_format=sten.CooTensor,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we create a new sparse model from the original dense model and run it with the same arguments as before:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/users/aivanov/sten1/src/sten/sten.py:1824: UserWarning: Sparse CSR tensor support is in beta state. If you miss a functionality in the sparse tensor support, please submit a feature request to https://github.com/pytorch/pytorch/issues. (Triggered internally at  /home/conda/feedstock_root/build_artifacts/pytorch-recipe_1664392091397/work/aten/src/ATen/SparseCsrTensorImpl.cpp:66.)\n",
      "  CsrTensor(tensor.to_sparse_csr()),\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([8, 128, 768])\n"
     ]
    }
   ],
   "source": [
    "sparse_model = sb.get_sparse_model(model)\n",
    "output = sparse_model(input)\n",
    "print(output[0].shape)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.13 ('sten1')",
   "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.9.13"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "faa923315ecf15f4653530526b0d3c584ef5f4f6244f0a40d7cbffbaf96f5375"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}