Jax benchmarks

benchmarks/gpu_results.csv

@@ -0,0 +1,25 @@
+,method,demeaner_backend,k,G,n_obs,full_feols_timing,demean_timing
+0,feols,jax,1,1,10000,0.22883999347686768,0.08213257789611816
+1,feols,jax,1,1,100000,0.27838222980499266,0.6599355697631836
+2,feols,jax,1,2,10000,0.2594131946563721,0.07008242607116699
+3,feols,jax,1,2,100000,0.5564646482467651,0.6762551069259644
+4,feols,jax,1,3,10000,0.29452004432678225,0.08089225292205811
+5,feols,jax,1,3,100000,0.7971727609634399,0.7361236095428467
+6,feols,jax,10,1,10000,0.23609719276428223,0.08923766613006592
+7,feols,jax,10,1,100000,0.35655946731567384,0.8534042596817016
+8,feols,jax,10,2,10000,0.2677977800369263,0.07926526069641113
+9,feols,jax,10,2,100000,0.7200102090835572,0.880342435836792
+10,feols,jax,10,3,10000,0.30958850383758546,0.0926206350326538
+11,feols,jax,10,3,100000,0.9418511629104614,0.8934823036193847
+12,feols,numba,1,1,10000,0.2558188199996948,0.0025653839111328125
+13,feols,numba,1,1,100000,0.26205954551696775,0.025234317779541014
+14,feols,numba,1,2,10000,0.23201637268066405,0.0021322011947631837
+15,feols,numba,1,2,100000,0.33342432975769043,0.03563899993896484
+16,feols,numba,1,3,10000,0.23092274665832518,0.0023543596267700194
+17,feols,numba,1,3,100000,0.25670347213745115,0.021407365798950195
+18,feols,numba,10,1,10000,0.4105666160583496,0.0046176910400390625
+19,feols,numba,10,1,100000,0.32753868103027345,0.08355832099914551
+20,feols,numba,10,2,10000,0.23588848114013672,0.005055141448974609
+21,feols,numba,10,2,100000,0.3476848602294922,0.0862342357635498
+22,feols,numba,10,3,10000,0.25297954082489016,0.005847930908203125
+23,feols,numba,10,3,100000,0.3176179170608521,0.056423068046569824

benchmarks/gpu_runtime_res.csv

@@ -0,0 +1,121 @@
+,method,demeaner_backend,k,G,n_obs,full_feols_timing,demean_timing
+0,feols,jax,1,1,1000,0.1320037841796875,
+1,feols,jax,1,1,10000,0.16295738220214845,
+2,feols,jax,1,1,100000,0.20828957557678224,
+3,feols,jax,1,1,1000000,0.42873706817626955,
+4,feols,jax,1,1,10000000,3.4382568836212157,
+5,feols,jax,1,2,1000,0.1420374870300293,
+6,feols,jax,1,2,10000,0.18153152465820313,
+7,feols,jax,1,2,100000,0.2269190788269043,
+8,feols,jax,1,2,1000000,0.45243015289306643,
+9,feols,jax,1,2,10000000,3.5151084899902343,
+10,feols,jax,1,3,1000,0.151939058303833,
+11,feols,jax,1,3,10000,0.18758950233459473,
+12,feols,jax,1,3,100000,0.2361727237701416,
+13,feols,jax,1,3,1000000,0.4868768692016602,
+14,feols,jax,1,3,10000000,3.8493817329406737,
+15,feols,jax,10,1,1000,0.19514913558959962,
+16,feols,jax,10,1,10000,0.17230439186096191,
+17,feols,jax,10,1,100000,0.2664950370788574,
+18,feols,jax,10,1,1000000,0.9469316005706787,
+19,feols,jax,10,1,10000000,8.387261295318604,
+20,feols,jax,10,2,1000,0.20477962493896484,
+21,feols,jax,10,2,10000,0.18998703956604004,
+22,feols,jax,10,2,100000,0.26862139701843263,
+23,feols,jax,10,2,1000000,0.9852738380432129,
+24,feols,jax,10,2,10000000,8.597736883163453,
+25,feols,jax,10,3,1000,0.19365372657775878,
+26,feols,jax,10,3,10000,0.2131319522857666,
+27,feols,jax,10,3,100000,0.28316402435302734,
+28,feols,jax,10,3,1000000,1.0177009105682373,
+29,feols,jax,10,3,10000000,8.918838596343994,
+30,feols,jax,50,1,1000,0.18034906387329103,
+31,feols,jax,50,1,10000,0.2439594268798828,
+32,feols,jax,50,1,100000,0.5509981155395508,
+33,feols,jax,50,1,1000000,3.538271760940552,
+34,feols,jax,50,1,10000000,33.43665189743042,
+35,feols,jax,50,2,1000,0.20095858573913575,
+36,feols,jax,50,2,10000,0.25174636840820314,
+37,feols,jax,50,2,100000,0.5164000511169433,
+38,feols,jax,50,2,1000000,3.4102453231811523,
+39,feols,jax,50,2,10000000,33.30270967483521,
+40,feols,jax,50,3,1000,0.20151491165161134,
+41,feols,jax,50,3,10000,0.26232194900512695,
+42,feols,jax,50,3,100000,0.529056167602539,
+43,feols,jax,50,3,1000000,3.44671630859375,
+44,feols,jax,50,3,10000000,33.779102373123166,
+45,feols,jax,100,1,1000,0.24306898117065429,
+46,feols,jax,100,1,10000,0.2974048137664795,
+47,feols,jax,100,1,100000,1.0358835697174071,
+48,feols,jax,100,1,1000000,6.931807231903076,
+49,feols,jax,100,1,10000000,68.10868678092956,
+50,feols,jax,100,2,1000,0.2459031581878662,
+51,feols,jax,100,2,10000,0.3029839515686035,
+52,feols,jax,100,2,100000,0.8705804347991943,
+53,feols,jax,100,2,1000000,6.833173513412476,
+54,feols,jax,100,2,10000000,67.79917821884155,
+55,feols,jax,100,3,1000,0.25403652191162107,
+56,feols,jax,100,3,10000,0.3005673885345459,
+57,feols,jax,100,3,100000,0.8948062419891357,
+58,feols,jax,100,3,1000000,6.897133874893188,
+59,feols,jax,100,3,10000000,68.098957157135,
+60,feols,numba,1,1,1000,0.1151197910308838,
+61,feols,numba,1,1,10000,0.11886963844299317,
+62,feols,numba,1,1,100000,0.17658085823059083,
+63,feols,numba,1,1,1000000,0.38565616607666015,
+64,feols,numba,1,1,10000000,3.7955574989318848,
+65,feols,numba,1,2,1000,0.13867654800415039,
+66,feols,numba,1,2,10000,0.12686729431152344,
+67,feols,numba,1,2,100000,0.18665246963500975,
+68,feols,numba,1,2,1000000,0.47369937896728515,
+69,feols,numba,1,2,10000000,5.122746992111206,
+70,feols,numba,1,3,1000,0.13739519119262694,
+71,feols,numba,1,3,10000,0.1281270980834961,
+72,feols,numba,1,3,100000,0.20229177474975585,
+73,feols,numba,1,3,1000000,0.4988919734954834,
+74,feols,numba,1,3,10000000,5.58708610534668,
+75,feols,numba,10,1,1000,0.34712915420532225,
+76,feols,numba,10,1,10000,0.12650370597839355,
+77,feols,numba,10,1,100000,0.20991711616516112,
+78,feols,numba,10,1,1000000,0.9024174690246582,
+79,feols,numba,10,1,10000000,9.295429277420045,
+80,feols,numba,10,2,1000,0.13671212196350097,
+81,feols,numba,10,2,10000,0.14796338081359864,
+82,feols,numba,10,2,100000,0.21911492347717285,
+83,feols,numba,10,2,1000000,1.0395352363586425,
+84,feols,numba,10,2,10000000,11.671601057052612,
+85,feols,numba,10,3,1000,0.144820499420166,
+86,feols,numba,10,3,10000,0.14316558837890625,
+87,feols,numba,10,3,100000,0.23515863418579103,
+88,feols,numba,10,3,1000000,1.0837657451629639,
+89,feols,numba,10,3,10000000,12.424165201187133,
+90,feols,numba,50,1,1000,0.13538317680358886,
+91,feols,numba,50,1,10000,0.19253530502319335,
+92,feols,numba,50,1,100000,0.47671823501586913,
+93,feols,numba,50,1,1000000,3.326345920562744,
+94,feols,numba,50,1,10000000,36.28279056549072,
+95,feols,numba,50,2,1000,0.15496983528137206,
+96,feols,numba,50,2,10000,0.2061081886291504,
+97,feols,numba,50,2,100000,0.5060088634490967,
+98,feols,numba,50,2,1000000,3.9241212368011475,
+99,feols,numba,50,2,10000000,46.66024560928345,
+100,feols,numba,50,3,1000,0.1508333683013916,
+101,feols,numba,50,3,10000,0.21372389793395996,
+102,feols,numba,50,3,100000,0.5096695899963379,
+103,feols,numba,50,3,1000000,4.05219841003418,
+104,feols,numba,50,3,10000000,48.755423641204835,
+105,feols,numba,100,1,1000,0.18344516754150392,
+106,feols,numba,100,1,10000,0.23093295097351074,
+107,feols,numba,100,1,100000,0.8638255596160889,
+108,feols,numba,100,1,1000000,6.876220750808716,
+109,feols,numba,100,1,10000000,84.45749440193177,
+110,feols,numba,100,2,1000,0.2027737617492676,
+111,feols,numba,100,2,10000,0.244942045211792,
+112,feols,numba,100,2,100000,0.9162676334381104,
+113,feols,numba,100,2,1000000,7.880603408813476,
+114,feols,numba,100,2,10000000,105.2484076499939,
+115,feols,numba,100,3,1000,0.1997222423553467,
+116,feols,numba,100,3,10000,0.2807606220245361,
+117,feols,numba,100,3,100000,0.9424882888793945,
+118,feols,numba,100,3,1000000,8.141235780715942,
+119,feols,numba,100,3,10000000,109.21670022010804,

docs/_quarto.yml

@@ -41,6 +41,8 @@ website:
           text: "Regression Decomposition"
         - text: "Compare fixest & PyFixest"
           file: compare-fixest-pyfixest.qmd
+        - text: "PyFixest on the GPU"
+          file: pyfixest_gpu.ipynb
         - text: "Replicating 'The Effect' with PyFixest"
           file: replicating-the-effect.qmd
 

docs/pyfixest_gpu.ipynb

@@ -0,0 +1,53 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## `PyFixest` on professional-tier GPUs \n",
+    "\n",
+    "`PyFixest` allows to run the fixed effects demeaning on the GPU via the `demeaner_backend` argument. \n",
+    "To do so, you will have to install `jax` and `jaxblib`, for example by typing `pip install pyfixest[jax]`.\n",
+    "\n",
+    "We test two back-ends for the iterative alternating-projections component of the fixed-effects regression on an Nvidia A100 GPU with 40 GB VRAM (a GPU that one typically wouldn't have installed to play graphics-intensive videogames on consumer hardware). `numba` benchmarks are run on a 12-core xeon CPU. \n",
+    "\n",
+    "The JAX backend exhibits major performance improvements **on the GPU** over numba in large problems. "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "![](figures/gpu_benchmarks.png)\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "On the **CPU** instead, we find that `numba` outperforms the JAX backend. You can find details in the [benchmark section](https://github.com/py-econometrics/pyfixest/tree/master/benchmarks) of the github repo. "
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "dev",
+   "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.12.8"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}