GPU Memory Exhaustion During Model Fitting on NVIDIA T4

Hi team,

I'm running into a GPU memory issue while fitting an MMM model and would appreciate some guidance on whether this appears to be a hardware limitation or a potential configuration issue.

**Model Configuration**

- Time periods: 66 weeks
- Geographies (DMAs): 183
- Media channels: 36
- Controls: 10 non-media variables + 8 controls

**Compute Environment**

Databricks cluster with:
- 4 × NVIDIA T4 GPUs
- 16 GB VRAM per GPU
- 440 GB system RAM

**Error**

During model fitting, I receive the following error:

Local rendezvous is aborting with status: RESOURCE_EXHAUSTED: OOM when allocating tensor with shape [36,7000,12078] and type float on/job:localhost/replica:0/task:0/device:GPU:0by allocator GPU_0_bfc

**My Interpretation**

The error appears to indicate that TensorFlow/JAX is attempting to allocate a tensor that exceeds the memory available on a single T4 GPU (GPU:0).

However, I'm unsure whether:

1. This dataset size is expected to exceed the practical limits of a 16 GB T4 GPU when using Meridian.
2. There may be a configuration issue causing excessive memory usage.
3. Meridian currently utilizes only a single GPU during this stage of model fitting, despite multiple GPUs being available on the cluster.

**Questions**

1. Based on the model dimensions (66 weeks, 183 DMAs, 36 media channels), would you expect a T4 with 16 GB VRAM to be sufficient?
2. Does the error indicate a genuine GPU memory limitation, or would you recommend investigating other aspects of the implementation/configuration?
3. If this is primarily a memory limitation, what would you recommend as the next best option:
- Reducing the number of geographies?
- Aggregating media channels?
- Reducing posterior samples/chains?
- Moving to a larger GPU (e.g., L4, A10, A100)?

Any guidance or recommendations would be greatly appreciated.

Thank you.

No repository files or tests are named. Start by reproducing model fitting with the stated dimensions and inspecting the reported tensor allocation and GPU usage across the four devices. Done means establishing whether the configuration exceeds one GPU's capacity or exposes a Meridian configuration or multi-GPU issue, with the supported limits or next steps documented.

Written by the indexing model from the issue text.