PySINDy with GPUs

[akaptano]

Megan and Michelle from Steve's lab are interested in adding GPU functionality to PySINDy as part of a class project. In principle, SINDy should be the easiest form of parallelization, since each x_dot = ..., y_dot = ..., etc. is separate (the coefficients being optimized are non-overlapping). Moreover, this would facilitate different SINDy libraries for each of these equations. Could imagine each GPU gets an equation, solves the miniature optimization problem, and combine at the end. We could also get speedups (with the in-house optimization algorithms, not sure we can do this with sklearn) because most of these rely on matrix-vector products and some form of least-squares, which can presumably be parallelized well on GPUs. Megan and Michelle could implement it for STLSQ for their class, and we could have undergrads extend it as needed.

@znicolaou @briandesilva Any thoughts on this or better ideas about how GPUs could be helpful?

[znicolaou]

Sounds like an good idea! Here are a few possibilities:

1) The biggest computational roadblock for me have been with spatiotemporal data, simply because the number of sample points gets large. Evaluation of the numerical derivatives and the application of the library functions across all space-time points seems like they could be implemented on a gpu, and it may be pretty straightforward with tensorflow, for example. It may be a pain to admit either tensor objects or numpy arrays everywhere, but maybe it won't be hard.

2) BLAS operations like linear solvers have also been implemented on gpus that could be used in the iterative optimizers like the ridge regression, but the benefits are smaller than for "embarrassingly parallelizable" functions, and there may be no benefit for smaller datasets.

3) In practice, the sweeping across optimizer hyperparameters turns out to be a common time sync for me, trying to find the Pareto front etc. Parallelizing across thresholds and regularization coefficients could be implemented on a gpu, and it may be helpful for productivity.

[briandesilva]

My mental model here is along the lines of "is there a problem pysindy users currently face which we can help solve with GPUs?" I want to make sure we're starting with a problem and trying to figure out the best solution rather than starting with a solution and looking for a problem to use it on.

I have a few concerns around adding GPU functionality. One fear is that we'd add extra complexity to the package (more dependencies, potentially more complicated code), making it harder for new users to adopt. It may also be difficult to test the new functionality as I don't know if github has GPUs you can use for automated testing. Finally, I worry that few users will have access to GPUs which will cause the new functionality to be underutilized.

On the other hand, the first problem Zack raised seems like a decent contender for speedup via GPU. I can also see how (3) can be slow, but that's relatively easy for users to parallelize themselves e.g. with ray-tune.

[Jacob-Stevens-Haas]

Interesting... I don't really understand the nuts and bolts of how GPUs accelerate code. I had thought it was because the size of processor cache was larger on a GPU, not a question of which parts of algorithms are "embarrassingly parallelizable."

I would imagine the greater speedup and easier problem to solve would be fitting separate regressions, computing spatial derivatives at different timepoints, and time derivatives at different spatial points using some form of async. OTOH, if cupy is a good drop in replacement for scipy and numpy, maybe that's easier? Still recognizing @briandesilva's comment about "is this a real problem", though.

For testing, do the AMATH servers have GPUs? If so, it should be possible to add an ssh secret to GitHub actions to queue a test run on UW servers, but that's a bit out of my depth.