Doubts about using more than one GPU

[Vardioratai]

Hello!

I am using Meridian to perform performance experiments and I have generated a synthetic dataset to experiment time differences between CPU and GPU.

The timings between using a configuration with CPU and GPU is very clear, my question goes about using more than 1 GPU. Here is an example:

• Experiment with 50 geos, 5 years of weekly history and 5 media vars + 3 controls.
• AWS g4dn.2xlarge machine with 1 NVIDIA T4 Tensor Core GPU needs 37 min to train a model.
• AWS g4dn.12xlarge machine with 4 GPUs needs 32 min to train the same model. I use JAX code to enable the use of the 4 GPUs, I don't know if the problem is here or it's something from Meridian and I'm going crazy.

Can you help me to solve this doubt?

[cpulavarthi]

Hello @Vardioratai,

Thank you for reaching out and sharing a detailed description of your issue!

I would like to inform you that using more than 1 GPU isn't currently supported with Meridian. We use tensorflow function graphs to use gpu acceleration in sampling posterior, and the sampling is done in one GPU.

Our team is constantly working on enhancing Meridian's functionality and multi-GPU acceleration shall be considered in future releases. You can track the previous changes and unreleased feature updates in our Change Log.

Do reach out if you have any questions or suggestions related to Meridian.

Google Meridian Support Team