math.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "57223509-6b5e-4079-810e-389060c044e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sympy import *"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "473a033e-39c4-4c0f-acc7-777c8d05d3ed",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3f6afe92-9ab1-4672-a4bd-dfb8aae34ad4",
   "metadata": {},
   "outputs": [],
   "source": [
    "var('A1 A2 A3 A4')\n",
    "A_cumsum = Matrix([[0, 0, 0, 0,0], \n",
    "                   [A1, 0, 0,0,0],\n",
    "                   [A2, A2-A1,0,0,0],\n",
    "                   [A3,A3-A1,A3-A2,0,0],\n",
    "                   [A4, A4-A1, A4-A2, A4-A3,0],\n",
    "                  ]).applyfunc(lambda x: exp(x) if x else 0)\n",
    "A_cumsum1 = Matrix([[0,0,0],\n",
    "                   [A1,0,0],\n",
    "                   [A2, A2-A1, 0]]).applyfunc(lambda x: exp(x) if x else 0)\n",
    "A_cumsum2 = Matrix([[0,0,0],\n",
    "                   [A3,0,0],\n",
    "                   [A4,A4-A3,0]]).applyfunc(lambda x: exp(x) if x else 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c8d6fea5-68a4-4d14-85f5-a71f940f2cf4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "⎡ 0       0          0          0      0⎤\n",
      "⎢                                       ⎥\n",
      "⎢ A₁                                    ⎥\n",
      "⎢ℯ        0          0          0      0⎥\n",
      "⎢                                       ⎥\n",
      "⎢ A₂   -A₁ + A₂                         ⎥\n",
      "⎢ℯ    ℯ              0          0      0⎥\n",
      "⎢                                       ⎥\n",
      "⎢ A₃   -A₁ + A₃   -A₂ + A₃              ⎥\n",
      "⎢ℯ    ℯ          ℯ              0      0⎥\n",
      "⎢                                       ⎥\n",
      "⎢ A₄   -A₁ + A₄   -A₂ + A₄   -A₃ + A₄   ⎥\n",
      "⎣ℯ    ℯ          ℯ          ℯ          0⎦\n",
      "⎡ 0       0      0⎤\n",
      "⎢                 ⎥\n",
      "⎢ A₁              ⎥\n",
      "⎢ℯ        0      0⎥\n",
      "⎢                 ⎥\n",
      "⎢ A₂   -A₁ + A₂   ⎥\n",
      "⎣ℯ    ℯ          0⎦\n",
      "⎡ 0       0      0⎤\n",
      "⎢                 ⎥\n",
      "⎢ A₃              ⎥\n",
      "⎢ℯ        0      0⎥\n",
      "⎢                 ⎥\n",
      "⎢ A₄   -A₃ + A₄   ⎥\n",
      "⎣ℯ    ℯ          0⎦\n"
     ]
    }
   ],
   "source": [
    "pprint(A_cumsum)\n",
    "pprint(A_cumsum1)\n",
    "pprint(A_cumsum2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "facea432-3114-416f-8a93-a849ccc0d225",
   "metadata": {},
   "outputs": [],
   "source": [
    "var('hi, hw, hx, hy, hz')\n",
    "states = Matrix([0, hw, hx, hy, hz])\n",
    "states1 = Matrix([0,hw,hx])\n",
    "states2 = Matrix([0,hy,hz])\n",
    "hi = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "293d9bc4-cb03-4f00-b5e2-6ccf053ffd5f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "⎡                    0                     ⎤\n",
      "⎢                                          ⎥\n",
      "⎢                    0                     ⎥\n",
      "⎢                                          ⎥\n",
      "⎢                   -A₁ + A₂               ⎥\n",
      "⎢               hw⋅ℯ                       ⎥\n",
      "⎢                                          ⎥\n",
      "⎢           -A₁ + A₃       -A₂ + A₃        ⎥\n",
      "⎢       hw⋅ℯ         + hx⋅ℯ                ⎥\n",
      "⎢                                          ⎥\n",
      "⎢    -A₁ + A₄       -A₂ + A₄       -A₃ + A₄⎥\n",
      "⎣hw⋅ℯ         + hx⋅ℯ         + hy⋅ℯ        ⎦\n",
      "⎡     0      ⎤\n",
      "⎢            ⎥\n",
      "⎢     0      ⎥\n",
      "⎢            ⎥\n",
      "⎢    -A₁ + A₂⎥\n",
      "⎣hw⋅ℯ        ⎦\n",
      "⎡     0      ⎤\n",
      "⎢            ⎥\n",
      "⎢     0      ⎥\n",
      "⎢            ⎥\n",
      "⎢    -A₃ + A₄⎥\n",
      "⎣hy⋅ℯ        ⎦\n"
     ]
    }
   ],
   "source": [
    "pprint(A_cumsum.multiply(states))\n",
    "pprint(A_cumsum1.multiply(states1))\n",
    "pprint(A_cumsum2.multiply(states2))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45c612b3-d242-43b1-87d0-54fd2329125b",
   "metadata": {},
   "source": [
    "# State passing test code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "aa8edd9e-7275-4f18-8fbd-9fc9417b4335",
   "metadata": {},
   "outputs": [],
   "source": [
    "#The reference function for state passing out of Mamba 2\n",
    "import torch.nn.functional as F\n",
    "def state_passing_ref(states, dA_chunk_cumsum, initial_states=None):\n",
    "    \"\"\"\n",
    "    Argument:\n",
    "        states: (batch, nchunks, nheads, dim)\n",
    "        dA_chunk_cumsum: (batch, nheads, nchunks)\n",
    "        initial_states: (batch, nheads, dim)\n",
    "    Return:\n",
    "        out: (batch, nchunks, nheads, dim)\n",
    "        final_states: (batch, nheads, dim)\n",
    "    \"\"\"\n",
    "    if initial_states is None:\n",
    "        initial_states = torch.zeros_like(states[:, 0])\n",
    "    states = torch.cat([rearrange(initial_states, \"b h d -> b 1 h d\"), states], dim=1)\n",
    "    dA_chunk_cumsum = F.pad(dA_chunk_cumsum, (1, 0))\n",
    "    dA_chunk_cumsum = torch.cumsum(dA_chunk_cumsum, dim=-1)\n",
    "    nchunks = dA_chunk_cumsum.shape[-1]\n",
    "    # (batch, nheads, nchunks, nchunks)\n",
    "    dt_chunk_segment_sum = dA_chunk_cumsum[:, :, :, None] - dA_chunk_cumsum[:, :, None, :]\n",
    "    # (batch, nheads, nchunks, nchunks)\n",
    "    decay_chunk = torch.exp(dt_chunk_segment_sum)\n",
    "    causal_mask = torch.tril(torch.ones(nchunks, nchunks, device=states.device, dtype=bool), diagonal=0)\n",
    "    decay_chunk = decay_chunk.masked_fill(~causal_mask, 0)\n",
    "    out = torch.einsum(\"bhzc,bchd->bzhd\", decay_chunk.to(dtype=states.dtype), states)\n",
    "    return out[:, :-1], out[:, -1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "7a701734-0532-49c6-b40d-22908b115584",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Trimmed it down to 1 batch 1 head\n",
    "def state_passing_test(states, dA_chunk_cumsum, initial_states=None):\n",
    "    \"\"\"\n",
    "    Argument:\n",
    "        states: (nchunks, dim)\n",
    "        dA_chunk_cumsum: (batch, nheads, nchunks)\n",
    "        initial_states: (batch, nheads, dim)\n",
    "    Return:\n",
    "        out: (batch, nchunks, nheads, dim)\n",
    "        final_states: (batch, nheads, dim)\n",
    "    \"\"\"\n",
    "    if initial_states is None:\n",
    "        initial_states = torch.zeros_like(states[0])\n",
    "        print(rearrange(initial_states, \"d -> 1 d\").shape)\n",
    "    #print(f\"{states.shape = }\")\n",
    "    states = torch.cat([rearrange(initial_states, \"d -> 1 d\"), states], dim=0)\n",
    "    dA_chunk_cumsum = F.pad(dA_chunk_cumsum, (1, 0))\n",
    "    dA_chunk_cumsum = torch.cumsum(dA_chunk_cumsum, dim=-1)\n",
    "    nchunks = dA_chunk_cumsum.shape[-1]\n",
    "    # (batch, nheads, nchunks, nchunks)\n",
    "    dt_chunk_segment_sum = dA_chunk_cumsum[:, None] - dA_chunk_cumsum[None, :]\n",
    "    #print(f\"{dt_chunk_segment_sum = }\")\n",
    "    # (batch, nheads, nchunks, nchunks)\n",
    "    #decay_chunk = torch.exp(dt_chunk_segment_sum)\n",
    "    decay_chunk=dt_chunk_segment_sum\n",
    "    #print(f\"{decay_chunk.shape = }\")\n",
    "    causal_mask = torch.tril(torch.ones(nchunks, nchunks, device=states.device, dtype=bool), diagonal=0)\n",
    "    #print(f\"{causal_mask.shape = }\")\n",
    "    decay_chunk = decay_chunk.masked_fill(~causal_mask, 0)\n",
    "    #print(f\"{decay_chunk = }\")\n",
    "    out = torch.einsum(\"zc,cd->zd\", decay_chunk.to(dtype=states.dtype), states)\n",
    "    return out[:-1], out[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "3fd89aa1-c2e1-411e-aa75-a78f3f93ef9d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 3])\n",
      "torch.Size([1, 3])\n",
      "tensor([[0., 0., 0.],\n",
      "        [0., 0., 0.],\n",
      "        [0., 0., 0.],\n",
      "        [1., 2., 3.],\n",
      "        [0., 0., 0.],\n",
      "        [3., 6., 9.]])\n",
      "tensor([[ 0.,  0.,  0.],\n",
      "        [ 0.,  0.,  0.],\n",
      "        [ 1.,  2.,  3.],\n",
      "        [ 3.,  6.,  9.],\n",
      "        [ 6., 12., 18.],\n",
      "        [10., 20., 30.]])\n"
     ]
    }
   ],
   "source": [
    "#Let's test what happens when we pass states and dA chunk wise, using the output as input to the next state passing.\n",
    "import torch\n",
    "from einops import rearrange\n",
    "import torch.nn.functional as F\n",
    "\n",
    "chunks=6\n",
    "states = torch.ones([chunks,3]).cumsum(1)\n",
    "dA_chunk_cumsum = torch.ones([chunks])#.cumsum(0)\n",
    "\n",
    "#Get the total sout_0 after state passing\n",
    "sout_0, fout_0 = state_passing_test(states,dA_chunk_cumsum)\n",
    "\n",
    "#Run chunk wise\n",
    "all_sout=[]\n",
    "all_fout=[]\n",
    "fout = None\n",
    "gchunk=2 #chunks per simulated gpu\n",
    "for i in range(0,chunks//gchunk):\n",
    "    j = i+1\n",
    "    #print(f\"{i*gchunk}:{j*gchunk}\")\n",
    "    sout,fout = state_passing_test(states[i*gchunk:j*gchunk],\n",
    "                                   dA_chunk_cumsum[i*gchunk:j*gchunk],\n",
    "                                  initial_states=fout)\n",
    "    all_sout.append(sout)\n",
    "    all_fout.append(fout)\n",
    "#print(fout)\n",
    "#print(fout_0)\n",
    "# These should be the same\n",
    "print(torch.cat(all_sout))\n",
    "print(sout_0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd55ca72-5e36-4399-8234-a608948ba566",
   "metadata": {},
   "source": [
    "## Now the goal for all-to-all\n",
    "- Given all_fout, all_sout, and dA_chunk_cumsum, where no initial state is used in any of the updates, update all_sout to be sout_0 without using any sharing of information between all_sout elements"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "d7b4d807-8e75-4e40-b588-9c087dfd75de",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 3])\n",
      "torch.Size([1, 3])\n",
      "torch.Size([1, 3])\n",
      "torch.Size([1, 3])\n",
      "tensor([[0., 0., 0.],\n",
      "        [0., 0., 0.],\n",
      "        [0., 0., 0.],\n",
      "        [0., 0., 0.],\n",
      "        [0., 0., 0.],\n",
      "        [0., 0., 0.]])\n",
      "tensor([[ 0.,  0.,  0.],\n",
      "        [ 0.,  0.,  0.],\n",
      "        [ 1.,  2.,  3.],\n",
      "        [ 3.,  6.,  9.],\n",
      "        [ 6., 12., 18.],\n",
      "        [10., 20., 30.]])\n"
     ]
    }
   ],
   "source": [
    "#Let's test what happens when we pass states and dA chunk wise, using the output as input to the next state passing.\n",
    "import torch\n",
    "from einops import rearrange\n",
    "import torch.nn.functional as F\n",
    "\n",
    "chunks=6\n",
    "states = torch.ones([chunks,3]).cumsum(1)\n",
    "dA_chunk_cumsum = torch.ones([chunks])#.cumsum(0)\n",
    "\n",
    "#Get the total sout_0 after state passing\n",
    "sout_0, fout_0 = state_passing_test(states,dA_chunk_cumsum)\n",
    "\n",
    "#Run chunk wise\n",
    "all_sout=[]\n",
    "all_fout=[]\n",
    "fout = None\n",
    "gchunk=2 #chunks per simulated gpu\n",
    "for i in range(0,chunks//gchunk):\n",
    "    j = i+1\n",
    "    #print(f\"{i*gchunk}:{j*gchunk}\")\n",
    "    sout,fout = state_passing_test(states[i*gchunk:j*gchunk],\n",
    "                                   dA_chunk_cumsum[i*gchunk:j*gchunk],\n",
    "                                  )\n",
    "    all_sout.append(sout)\n",
    "    all_fout.append(fout)\n",
    "#print(fout)\n",
    "#print(fout_0)\n",
    "\n",
    "# Need to find a way to make these the same after sate_passing\n",
    "print(torch.cat(all_sout))\n",
    "print(sout_0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "377eb2ac-89dc-476b-9b9d-bd3f788b42ab",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "mamba",
   "language": "python",
   "name": "mamba"
  },
  "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.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}