Skip to content

Commit 2425240

Browse files
s3alfiscapoorvalal
s3alfisc
and
apoorvalal
authored
Jax benchmarks (#784)
* add jax benchmark notebook * address review comments * burnin for demean * add big gpu benchmarks * some cleanups & gpu notebook to docs * update lock file --------- Co-authored-by: Apoorva Lal <[email protected]>
1 parent 4a4bc1c commit 2425240

9 files changed

+3182
-1
lines changed

‎benchmarks/gpu_benchmarks.ipynb

+1,494
Large diffs are not rendered by default.

‎benchmarks/gpu_big_benchmarks.ipynb

+1,465
Large diffs are not rendered by default.

‎benchmarks/gpu_results.csv

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

‎benchmarks/gpu_runtime_res.csv

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

‎docs/_quarto.yml

+2
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,8 @@ website:
4141
text: "Regression Decomposition"
4242
- text: "Compare fixest & PyFixest"
4343
file: compare-fixest-pyfixest.qmd
44+
- text: "PyFixest on the GPU"
45+
file: pyfixest_gpu.ipynb
4446
- text: "Replicating 'The Effect' with PyFixest"
4547
file: replicating-the-effect.qmd
4648

‎docs/figures/gpu_benchmarks.png

78.4 KB
Loading

‎docs/pyfixest_gpu.ipynb

+53
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "markdown",
5+
"metadata": {},
6+
"source": [
7+
"## `PyFixest` on professional-tier GPUs \n",
8+
"\n",
9+
"`PyFixest` allows to run the fixed effects demeaning on the GPU via the `demeaner_backend` argument. \n",
10+
"To do so, you will have to install `jax` and `jaxblib`, for example by typing `pip install pyfixest[jax]`.\n",
11+
"\n",
12+
"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",
13+
"\n",
14+
"The JAX backend exhibits major performance improvements **on the GPU** over numba in large problems. "
15+
]
16+
},
17+
{
18+
"cell_type": "markdown",
19+
"metadata": {},
20+
"source": [
21+
"![](figures/gpu_benchmarks.png)\n"
22+
]
23+
},
24+
{
25+
"cell_type": "markdown",
26+
"metadata": {},
27+
"source": [
28+
"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. "
29+
]
30+
}
31+
],
32+
"metadata": {
33+
"kernelspec": {
34+
"display_name": "dev",
35+
"language": "python",
36+
"name": "python3"
37+
},
38+
"language_info": {
39+
"codemirror_mode": {
40+
"name": "ipython",
41+
"version": 3
42+
},
43+
"file_extension": ".py",
44+
"mimetype": "text/x-python",
45+
"name": "python",
46+
"nbconvert_exporter": "python",
47+
"pygments_lexer": "ipython3",
48+
"version": "3.12.8"
49+
}
50+
},
51+
"nbformat": 4,
52+
"nbformat_minor": 2
53+
}

0 commit comments

Comments
 (0)