pyfixest is MUCH slower than (Stata + Julia) reghdfejl

[rhstanton]

First, thanks very much for this package. It gets me a lot closer to being able to do everything in Python! However, I've noticed an enormous speed difference between pyfixest and the Stata (+ Julia) function reghdfejl. Specifically, using the exact same data and running the same regression with two sets of fixed effects, pyfixest takes 3 - 6 hours (depending on what I set the tolerance to), while reghdfejl takes something like 3 minutes. That's quite a difference! The results generated are very similar in both cases.

I'm using pyfixest 0.30.2 with jax 0.7.2 and cuda 13. pyfixest is definitely using the GPU (Nvidia 3090Ti), which shows 100% utilization during the entire run, has something over 18GB RAM allocated, and heats up my office nicely. reghdfejl is a Stata routine that calls Julia to do the heavy processing, and again I've set it to use the GPU.

I've seen older discussions about the speed not being what it might be, but this difference seems quite extreme. Any thoughts on what might be going on? One possibly relevant thought: I have to increase the maximum number of iterations for the demeaning to finish, otherwise it gets to 100,000 and stops with an error. That seems like quite a lot of iterations.

Thanks for any suggestions.

[s3alfisc]

Hi @rhstanton , thanks so much for your feedback! Very cool to learn that there are real users of the JAX backend! 3min vs 6h sounds indeed off and is something we need to tackle. Either it's a problem with the algorithm in itself (which I believe can likely be improved), but here it seems like a problem with non-convergence.

Non-Convergence

Are you fitting a problem with fixed effects that are highly nested? If so, this might already explain the difference. In this case, the pyfixest demeaning algorithm might struggle to converge relative to the Julia algorithm. We (@schroedk) have a WIP implementation of an acceleration scheme that helps in this case for the Rust backend here and should eventually also implement it for the JAX backend.

JAX Optimizations

@juanitorduz implemented the JAX backend back in January and we did some benchmarks, but we never spend any extra efforts to fully optimize the algo. For this, I personally lack experience with JAX plus also access to GPUs. But we should certainly try to improve here as well.

[rhstanton]

Thanks for the reply. There are two sets of fixed effects: property_id and year^fire_id. Each value for fire_id corresponds to a particular subset of the property_ids.

[s3alfisc]

Thanks, and sorry for the late response! I was struggling with aligning pyfixest with recent API changes in fixest plus all sorts of HAC SEs in the last days =) In this case it is likely the algorithm is slowly converging and not its JAX implementation per se. Will put emphasis on optimizing this after the next release (mid November).

[rhstanton]

Thanks! Meanwhile, I've added an option to my code so it can call FixedEffectModels.jl instead. Alexander Fischer ***@***.***> writes:

…

s3alfisc left a comment (py-econometrics/pyfixest#1042) Thanks, and sorry for the late response! I was struggling with aligning pyfixest with recent API changes in fixest plus all sorts of HAC SEs in the last days =) In this case it is likely the algorithm in JAX that is slowly converging. Will put emphasis on optimizing this after the next release (mid November). - Reply to this email directly, view it on GitHub, or unsubscribe. You are receiving this because you were mentioned.*Message ID: ***@***.***>

[s3alfisc]

Great! Out of curiosity, if you run the problem via Julia on the GPU vs CPU, how much time do you save?

[rhstanton]

On my current problem, it actually takes longer (wall time) with the GPU activated, though I can see that the demeaning time is *much* faster. I assume there's a lot more overhead involved with setting up the GPU, transferring data, etc. On another problem I was running a while ago via Stata (but still using Julia to do the work), I remember that one job took 20 minutes with GPU off versus 2 minutes with GPU on (code otherwise identical). That's a significant time difference, especially if you're sitting in front of the screen waiting... Alexander Fischer ***@***.***> writes:

…

s3alfisc left a comment (py-econometrics/pyfixest#1042) Great! Out of curiosity, if you run the problem via Julia on the GPU vs CPU, how much time do you save? - Reply to this email directly, view it on GitHub, or unsubscribe. You are receiving this because you were mentioned.*Message ID: ***@***.***>

[rhstanton]

Richard Stanton ***@***.***> writes:

…

On my current problem, it actually takes longer (wall time) with the GPU activated, though I can see that the demeaning time is *much* faster. I assume there's a lot more overhead involved with setting up the GPU, transferring data, etc. On another problem I was running a while ago via Stata (but still using Julia to do the work), I remember that one job took 20 minutes with GPU off versus 2 minutes with GPU on (code otherwise identical). That's a significant time difference, especially if you're sitting in front of the screen waiting... Alexander Fischer ***@***.***> writes: > s3alfisc left a comment (py-econometrics/pyfixest#1042) > > Great! Out of curiosity, if you run the problem via Julia on the GPU > vs CPU, how much time do you save? > > - > Reply to this email directly, view it on GitHub, or unsubscribe. > You are receiving this because you were mentioned.*Message ID: > ***@***.***>

[rhstanton]

Actually, the GPU was a bit faster with Julia, but the difference was small. Here are comparative timings on the same problem: Julia (with GPU): 1 minute 45 seconds Julia (without GPU): 2 minutes 5 seconds pyfixest (with GPU): 161 minutes pyfixest (using Rust): 197 minutes In both cases with pyfixest, I have fixef_tol set to 1e-06. My memory is that the time approximately doubles if I use the default value instead. Richard Stanton ***@***.***> writes:

…

On my current problem, it actually takes longer (wall time) with the GPU activated, though I can see that the demeaning time is *much* faster. I assume there's a lot more overhead involved with setting up the GPU, transferring data, etc. On another problem I was running a while ago via Stata (but still using Julia to do the work), I remember that one job took 20 minutes with GPU off versus 2 minutes with GPU on (code otherwise identical). That's a significant time difference, especially if you're sitting in front of the screen waiting... Alexander Fischer ***@***.***> writes: > s3alfisc left a comment (py-econometrics/pyfixest#1042) > > Great! Out of curiosity, if you run the problem via Julia on the GPU > vs CPU, how much time do you save? > > - > Reply to this email directly, view it on GitHub, or unsubscribe. > You are receiving this because you were mentioned.*Message ID: > ***@***.***>

[s3alfisc]

Uff, 2 min vs 2h is not something I'm happy with. I'll invest some time this weekend on the acceleration in Rust & will try to port it to JAX too. Likely I should also try to fully vecroize the JAX algos and trade off some memory efficiency for speed.

Would you be up to help me out with some test runs @rhstanton? I unfortunately don't have access to GPUs / not sure if I will get JAX to run on the M chips I have access to.

[rhstanton]

I'd be very happy to do some test runs. Just let me know. Best, Richard Alexander Fischer ***@***.***> writes:

…

s3alfisc left a comment (py-econometrics/pyfixest#1042) Uff, 2 min vs 2h is not something I'm happy with. I'll invest some time this weekend on the acceleration in Rust & will try to port it to JAX too. Likely I should also try to fully vecroize the JAX algos and trade off some memory efficiency for speed. Would you be up to help me out with some test runs @rhstanton? I unfortunately don't have access to GPUs / not sure if I will get JAX to run on the M chips I have access to. - Reply to this email directly, view it on GitHub, or unsubscribe. You are receiving this because you were mentioned.*Message ID: ***@***.***>

[rhstanton]

By the way, the Julia algorithm reports 181 iterations to get convergence. I don't know how many pyfixest does, but I know that at the default setting (maximum = 100,000) it dies complaining about too many iterations. I set it to 500,000 and it runs to completion. Of course, the total time also depends on what's being done during each iteration, but this difference is rather large... Alexander Fischer ***@***.***> writes:

…

s3alfisc left a comment (py-econometrics/pyfixest#1042) Uff, 2 min vs 2h is not something I'm happy with. I'll invest some time this weekend on the acceleration in Rust & will try to port it to JAX too. Likely I should also try to fully vecroize the JAX algos and trade off some memory efficiency for speed. Would you be up to help me out with some test runs @rhstanton? I unfortunately don't have access to GPUs / not sure if I will get JAX to run on the M chips I have access to. - Reply to this email directly, view it on GitHub, or unsubscribe. You are receiving this because you were mentioned.*Message ID: ***@***.***>

[s3alfisc]

I set it to 500,000 and it runs to completion.

Very interesting! I am fairly certain that this is due to the specific structure of your fixed effects problem, where the standard algo simply struggles to converge.

If you want to try bleeding edge pyfixest and a potentially faster algo (via the Irons-Tuck acceleration), you can install from this gh branch

https://github.com/schroedk/pyfixest/tree/feature/demean-accelerated

and then launch the dev environment (after installing pixi)

pixi install -e dev

which downloads/installs all dependencies in the dev env (not global), including Rust & Maturin.

Then you should be able to compile Rust by typing

pixi run -e dev maturin-develop

and

then you can run any python script in the dev env via

pixi r -e dev python_script.py

For example, the benchmarks for the accelerated algo:

pixi run -e dev pytest  tests/test_demean.py::test_demean_complex_fixed_effects

For N = 100K we get a 6-7x improvement:

(All of this works out of the box thanks to pixi).

And if your super curious to see how it works, I have sketched out the Irons Tucks details here plus a basic numba implementation - eventually I'd love to port this to JAX as well =)

Btw, I am not 100% sure what FixedEffectsModels.jl does under the hood - as far as I know, r-fixest uses the IT acceleration. Plus there are some fancier graph based methods described in the reghdfe papers, though I never looked at them too closely.

[s3alfisc]

Also, a long shot (I have to try, sry😅) - in case you'd be interested to try yourself at the acceleration for JAX, that'd be a fantastic contribution =)

[rhstanton]

Unlikely - I've never used Jax, except indirectly. But once everything else is stable, let me know and let's discuss what it would take. Alexander Fischer ***@***.***> writes:

…

s3alfisc left a comment (py-econometrics/pyfixest#1042) Also, a long shot (I have to try, sry😅) - in case you'd be interested to try yourself at the acceleration for JAX, that'd be a fantastic contribution =) - Reply to this email directly, view it on GitHub, or unsubscribe. You are receiving this because you were mentioned.*Message ID: ***@***.***>

[rhstanton]

Speed is not the only difference. When I run the Julia code, my GPU reports memory usage as 876MB. When I use pyfixest, it reports something over 18 *GB*. This is not 100% fair, as I think Jax grabs more than it actually needs, but it should be noted that before I converted all my numbers to 32 bit, the code crashed with an out-of-memory error. Richard Stanton ***@***.***> writes:

…

Actually, the GPU was a bit faster with Julia, but the difference was small. Here are comparative timings on the same problem: Julia (with GPU): 1 minute 45 seconds Julia (without GPU): 2 minutes 5 seconds pyfixest (with GPU): 161 minutes pyfixest (using Rust): 197 minutes In both cases with pyfixest, I have fixef_tol set to 1e-06. My memory is that the time approximately doubles if I use the default value instead. Richard Stanton ***@***.***> writes: > On my current problem, it actually takes longer (wall time) with > the GPU activated, though I can see that the demeaning time is > *much* faster. I assume there's a lot more overhead involved with > setting up the GPU, transferring data, etc. > > On another problem I was running a while ago via Stata (but still > using Julia to do the work), I remember that one job took 20 > minutes with GPU off versus 2 minutes with GPU on (code otherwise > identical). That's a significant time difference, especially if > you're sitting in front of the screen waiting... > > > > Alexander Fischer ***@***.***> writes: > >> s3alfisc left a comment (py-econometrics/pyfixest#1042) >> >> Great! Out of curiosity, if you run the problem via Julia on the GPU >> vs CPU, how much time do you save? >> >> - >> Reply to this email directly, view it on GitHub, or unsubscribe. >> You are receiving this because you were mentioned.*Message ID: >> ***@***.***>