speed.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Encoder Speed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import torch as t\n",
    "import torch.nn as nn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 1024\n",
    "input_dim = 512\n",
    "dict_size = 1_000_000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = t.randn(batch_size, input_dim).cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "W_small = nn.Linear(input_dim, dict_size // 16).cuda()\n",
    "W_big = nn.Linear(input_dim, dict_size).cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Forward pass took 3.243 milliseconds\n"
     ]
    }
   ],
   "source": [
    "# Warm-up run\n",
    "_ = W_small(x)\n",
    "\n",
    "# Ensure all CUDA operations are completed\n",
    "t.cuda.synchronize()\n",
    "\n",
    "# Time the operation\n",
    "start = t.cuda.Event(enable_timing=True)\n",
    "end = t.cuda.Event(enable_timing=True)\n",
    "\n",
    "start.record()\n",
    "output = W_small(x)\n",
    "end.record()\n",
    "\n",
    "# Ensure all CUDA operations are completed\n",
    "t.cuda.synchronize()\n",
    "\n",
    "elapsed_time = start.elapsed_time(end)\n",
    "print(f\"Forward pass took {elapsed_time:.3f} milliseconds\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Forward pass took 55.207 milliseconds\n"
     ]
    }
   ],
   "source": [
    "# Warm-up run\n",
    "_ = W_big(x)\n",
    "\n",
    "# Ensure all CUDA operations are completed\n",
    "t.cuda.synchronize()\n",
    "\n",
    "# Time the operation\n",
    "start = t.cuda.Event(enable_timing=True)\n",
    "end = t.cuda.Event(enable_timing=True)\n",
    "\n",
    "start.record()\n",
    "output = W_big(x)\n",
    "end.record()\n",
    "\n",
    "# Ensure all CUDA operations are completed\n",
    "t.cuda.synchronize()\n",
    "\n",
    "elapsed_time = start.elapsed_time(end)\n",
    "print(f\"Forward pass took {elapsed_time:.3f} milliseconds\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## MoE Library"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch as t\n",
    "from torch import nn\n",
    "from mixture_of_experts import MoE\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "moe1 = MoE(\n",
    "    dim = 768,\n",
    "    num_experts = 1,               # increase the experts (# parameters) of your model without increasing computation\n",
    "    hidden_dim = 768 * 32,           # size of hidden dimension in each expert, defaults to 4 * dimension\n",
    "    activation = nn.ReLU,      # use your preferred activation, will default to GELU\n",
    "    second_policy_train = 'random', # in top_2 gating, policy for whether to use a second-place expert\n",
    "    second_policy_eval = 'random',  # all (always) | none (never) | threshold (if gate value > the given threshold) | random (if gate value > threshold * random_uniform(0, 1))\n",
    "    second_threshold_train = 0.2,\n",
    "    second_threshold_eval = 0.2,\n",
    "    capacity_factor_train = 1.25,   # experts have fixed capacity per batch. we need some extra capacity in case gating is not perfectly balanced.\n",
    "    capacity_factor_eval = 2.,      # capacity_factor_* should be set to a value >=1\n",
    "    loss_coef = 1e-2                # multiplier on the auxiliary expert balancing auxiliary loss\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "moe4 = MoE(\n",
    "    dim = 768,\n",
    "    num_experts = 4,               # increase the experts (# parameters) of your model without increasing computation\n",
    "    hidden_dim = 768 * 32,           # size of hidden dimension in each expert, defaults to 4 * dimension\n",
    "    activation = nn.ReLU,      # use your preferred activation, will default to GELU\n",
    "    second_policy_train = 'random', # in top_2 gating, policy for whether to use a second-place expert\n",
    "    second_policy_eval = 'random',  # all (always) | none (never) | threshold (if gate value > the given threshold) | random (if gate value > threshold * random_uniform(0, 1))\n",
    "    second_threshold_train = 0.2,\n",
    "    second_threshold_eval = 0.2,\n",
    "    capacity_factor_train = 1.25,   # experts have fixed capacity per batch. we need some extra capacity in case gating is not perfectly balanced.\n",
    "    capacity_factor_eval = 2.,      # capacity_factor_* should be set to a value >=1\n",
    "    loss_coef = 1e-2                # multiplier on the auxiliary expert balancing auxiliary loss\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "moe16 = MoE(\n",
    "    dim = 768,\n",
    "    num_experts = 16,               # increase the experts (# parameters) of your model without increasing computation\n",
    "    hidden_dim = 768 * 32,           # size of hidden dimension in each expert, defaults to 4 * dimension\n",
    "    activation = nn.ReLU,      # use your preferred activation, will default to GELU\n",
    "    second_policy_train = 'random', # in top_2 gating, policy for whether to use a second-place expert\n",
    "    second_policy_eval = 'random',  # all (always) | none (never) | threshold (if gate value > the given threshold) | random (if gate value > threshold * random_uniform(0, 1))\n",
    "    second_threshold_train = 0.2,\n",
    "    second_threshold_eval = 0.2,\n",
    "    capacity_factor_train = 1.25,   # experts have fixed capacity per batch. we need some extra capacity in case gating is not perfectly balanced.\n",
    "    capacity_factor_eval = 2.,      # capacity_factor_* should be set to a value >=1\n",
    "    loss_coef = 1e-2                # multiplier on the auxiliary expert balancing auxiliary loss\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "moe64 = MoE(\n",
    "    dim = 768,\n",
    "    num_experts = 64,               # increase the experts (# parameters) of your model without increasing computation\n",
    "    hidden_dim = 768 * 32,           # size of hidden dimension in each expert, defaults to 4 * dimension\n",
    "    activation = nn.ReLU,      # use your preferred activation, will default to GELU\n",
    "    second_policy_train = 'random', # in top_2 gating, policy for whether to use a second-place expert\n",
    "    second_policy_eval = 'random',  # all (always) | none (never) | threshold (if gate value > the given threshold) | random (if gate value > threshold * random_uniform(0, 1))\n",
    "    second_threshold_train = 0.2,\n",
    "    second_threshold_eval = 0.2,\n",
    "    capacity_factor_train = 1.25,   # experts have fixed capacity per batch. we need some extra capacity in case gating is not perfectly balanced.\n",
    "    capacity_factor_eval = 2.,      # capacity_factor_* should be set to a value >=1\n",
    "    loss_coef = 1e-2                # multiplier on the auxiliary expert balancing auxiliary loss\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "moe1024 = MoE(\n",
    "    dim = 768,\n",
    "    num_experts = 1024,               # increase the experts (# parameters) of your model without increasing computation\n",
    "    hidden_dim = 768 * 32,           # size of hidden dimension in each expert, defaults to 4 * dimension\n",
    "    activation = nn.ReLU,      # use your preferred activation, will default to GELU\n",
    "    second_policy_train = 'random', # in top_2 gating, policy for whether to use a second-place expert\n",
    "    second_policy_eval = 'random',  # all (always) | none (never) | threshold (if gate value > the given threshold) | random (if gate value > threshold * random_uniform(0, 1))\n",
    "    second_threshold_train = 0.2,\n",
    "    second_threshold_eval = 0.2,\n",
    "    capacity_factor_train = 1.25,   # experts have fixed capacity per batch. we need some extra capacity in case gating is not perfectly balanced.\n",
    "    capacity_factor_eval = 2.,      # capacity_factor_* should be set to a value >=1\n",
    "    loss_coef = 1e-2                # multiplier on the auxiliary expert balancing auxiliary loss\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## are the weights stored on CPU?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs = t.randn(1, 8192, 768)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Forward pass took 3704.199 milliseconds\n"
     ]
    }
   ],
   "source": [
    "# Warm-up run\n",
    "_ = moe1(inputs)\n",
    "\n",
    "# Ensure all CUDA operations are completed\n",
    "t.cuda.synchronize()\n",
    "\n",
    "# Time the operation\n",
    "start = t.cuda.Event(enable_timing=True)\n",
    "end = t.cuda.Event(enable_timing=True)\n",
    "\n",
    "start.record()\n",
    "out, aux_loss = moe1(inputs)\n",
    "end.record()\n",
    "\n",
    "# Ensure all CUDA operations are completed\n",
    "t.cuda.synchronize()\n",
    "\n",
    "elapsed_time = start.elapsed_time(end)\n",
    "print(f\"Forward pass took {elapsed_time:.3f} milliseconds\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Forward pass took 5826.192 milliseconds\n"
     ]
    }
   ],
   "source": [
    "# Warm-up run\n",
    "_ = moe4(inputs)\n",
    "\n",
    "# Ensure all CUDA operations are completed\n",
    "t.cuda.synchronize()\n",
    "\n",
    "# Time the operation\n",
    "start = t.cuda.Event(enable_timing=True)\n",
    "end = t.cuda.Event(enable_timing=True)\n",
    "\n",
    "start.record()\n",
    "out, aux_loss = moe4(inputs)\n",
    "end.record()\n",
    "\n",
    "# Ensure all CUDA operations are completed\n",
    "t.cuda.synchronize()\n",
    "\n",
    "elapsed_time = start.elapsed_time(end)\n",
    "print(f\"Forward pass took {elapsed_time:.3f} milliseconds\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Forward pass took 8306.871 milliseconds\n"
     ]
    }
   ],
   "source": [
    "# Warm-up run\n",
    "_ = moe16(inputs)\n",
    "\n",
    "# Ensure all CUDA operations are completed\n",
    "t.cuda.synchronize()\n",
    "\n",
    "# Time the operation\n",
    "start = t.cuda.Event(enable_timing=True)\n",
    "end = t.cuda.Event(enable_timing=True)\n",
    "\n",
    "start.record()\n",
    "out, aux_loss = moe16(inputs)\n",
    "end.record()\n",
    "\n",
    "# Ensure all CUDA operations are completed\n",
    "t.cuda.synchronize()\n",
    "\n",
    "elapsed_time = start.elapsed_time(end)\n",
    "print(f\"Forward pass took {elapsed_time:.3f} milliseconds\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Forward pass took 7080.979 milliseconds\n"
     ]
    }
   ],
   "source": [
    "# Warm-up run\n",
    "_ = moe64(inputs)\n",
    "\n",
    "# Ensure all CUDA operations are completed\n",
    "t.cuda.synchronize()\n",
    "\n",
    "# Time the operation\n",
    "start = t.cuda.Event(enable_timing=True)\n",
    "end = t.cuda.Event(enable_timing=True)\n",
    "\n",
    "start.record()\n",
    "out, aux_loss = moe64(inputs)\n",
    "end.record()\n",
    "\n",
    "# Ensure all CUDA operations are completed\n",
    "t.cuda.synchronize()\n",
    "\n",
    "elapsed_time = start.elapsed_time(end)\n",
    "print(f\"Forward pass took {elapsed_time:.3f} milliseconds\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Warm-up run\n",
    "_ = moe1024(inputs)\n",
    "\n",
    "# Ensure all CUDA operations are completed\n",
    "t.cuda.synchronize()\n",
    "\n",
    "# Time the operation\n",
    "start = t.cuda.Event(enable_timing=True)\n",
    "end = t.cuda.Event(enable_timing=True)\n",
    "\n",
    "start.record()\n",
    "out, aux_loss = moe1024(inputs)\n",
    "end.record()\n",
    "\n",
    "# Ensure all CUDA operations are completed\n",
    "t.cuda.synchronize()\n",
    "\n",
    "elapsed_time = start.elapsed_time(end)\n",
    "print(f\"Forward pass took {elapsed_time:.3f} milliseconds\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}