`parallel=True` does not fully reach `QuasisepSolver.condition`, making conditioning ~77x slower on GPU than serial on CPU
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
-
The
X_test=Nonebranch ofQuasisepSolver.conditionignoresself.parallel:
https://github.com/dfm/tinygp/blob/5302d5a1ae47ad9ff3fe857eee4cbbebbc1c5637/src/tinygp/solvers/quasisep/solver.py#L125-L129
QSM.__matmul__(core.py) has no keyword to pass one and dispatches toops.qsm_mul, andgram()isself.transpose() @ self, so three total sequential scans occur. -
gp.conditionconstructs the returnedGaussianProcesswithout forwarding solver kwargs, socond.solver.parallelisFalseeven whengp.solver.parallelisTrue, affecting all subsequent calls with the conditioned GP object.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- 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