dfm / dfm/tinygp

`parallel=True` does not fully reach `QuasisepSolver.condition`, making conditioning ~77x slower on GPU than serial on CPU

Open
#280 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
349
Forks
35
Avg merge
2d 2h
Merged PRs (30d)
1

Description

PR #269 works as advertised for the vector-facing operations, and log_probability with parallel=True is 322x faster on GPU than on CPU (though it performs about the same as parallel=False on CPU).

However, the new parallel methods only partially reach condition, and the conditioned GP that comes back doesn't inherit the parallel flag.

I noticed this when benchmarking this new functionality against smolgp. See on that page the benchmarking plots for the likelihood and conditioning calls, for instantaneous data.

Environment
  • tinygp 0.3.2.dev13+g5302d5a1a (includes #269)
  • jax / jaxlib 0.9.2, float64 enabled
  • GPU: NVIDIA RTX 6000 Ada, 48 GB · CPU: Intel Xeon w5-3435X (32 cores)
  • Kernel: tinygp.kernels.quasisep.SHO (J = 2), sorted 1-D inputs
Measurements

Median of 7 repeats, each in a fresh process, timed around block_until_ready() and excluding JIT warm-up. Seconds.

N log_probability condition
serial CPU serial GPU parallel GPU serial CPU parallel GPU
5.6e4 0.009 1.57 0.010 0.055 3.15
3.2e5 0.024 8.96 0.042 0.261 18.06
1.8e6 0.114 52.52 0.211 1.432 105.35
1.0e7 0.585 297.16 0.922 7.980 613.68

Issue

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with src/tinygp/solvers/quasisep/solver.py at the X_test=None branch of QuasisepSolver.condition, then inspect QSM.matmul and ops.qsm_mul in core.py. Trace gp.condition to see how the conditioned GaussianProcess is constructed. Done means conditioning honors the solver's parallel setting and the returned conditioned GP preserves it for later calls; compare against the issue's conditioning benchmarks.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning, performance
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Quiet
Clarity
Clearly specified
Newbie friendliness
68/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.