State Space: Anyone have a working colab config?

Questions version agnostic
,
1

Hey PyMC:

I’ve been trying colab as a way to improve modelling speed. Does anyone have a working config for colab for state space? (Edit: I’m using a model on colab that I’ve successfully run locally.)

I tried using default nuts sampler and numpyro_nuts and was able to get sampling. But the tree depth goes to 255 after the first 10 samples and then pretty much grinds to a halt. Maybe initial values aren’t being properly set?

I could not get nutpie installed. I got an error about a subprocess failing - and Gemini suggested nutpie needed Rust to compile. Which brought me to conda and the pre-compiled libraries. I also tried to install a conda environment - but was somewhat stymied in the attempt, but I forget why.

I am using the jax backend and set jax to cuda. I used the A100 GPU, and I tried a few others.

Anyway, any help you could provide to this first-timer would be appreciated.

Thank you for your time,
Roy

2

This seems to be more of a nutpie question than a statespace question. @aseyboldt would know if nutpie can get up and running on colab or not.

Regarding GPUs, they’re aren’t magical speedup machines. Only specific types of models show large speedups – basically those that involve large matrix multiplications in parallel (e.g. “wide” computation). Statespace models are “long” computation, doing lots of small matrix multiplications in sequence. I have not seen any benefit to GPU sampling with them.

1 Like
3

Great to know about the GPU.

I’ll leave the thread open for a few more days to see if anyone has a config they’ve used with the state space module that works really well on Colab or another platform. I’d like to know if I can use a platform before I buy a better CPU. (-:

Thank you, Jesse.

-Roy