patrick-kidger / patrick-kidger/diffrax
Question: DirectAdjoint is faster than RecursiveCheckpointAdjoint?
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 2.1k
- Forks
- 189
- Avg merge
- 3d 18h
- Merged PRs (30d)
- 1
Description
Hi,
According to the suggestions in adjoints docs, the RecursiveCheckpointAdjoint method, given enough checkpoints, should be faster than DirectAdjoint. But In my practice, it turns out that DirectAdjoint is faster. Is there wrong in my understanding?
Example
An example can be shown using the neural_cde tutorial:
- If use the default setting:
adjoint = RecursiveCheckpointAdjoint(), the total run time (without the first step/compilation time ) is 1.38s - if use
adjoint = RecursiveCheckpointAdjoint(checkpoints=4096), the run time is 1.18s - if use 'adjoint = diffrax.DirectAdjoint()', the run time is 0.77s
Environment
jax == 0.4.29jaxlib == 0.4.29diffrax == 0.6.0- platform: CPU
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
Read the adjoints documentation and reproduce the timings from examples/neural_cde.ipynb with the stated JAX, jaxlib, diffrax, and CPU environment. Compare RecursiveCheckpointAdjoint with different checkpoint counts against DirectAdjoint; done means explaining the observed performance difference or identifying a documentation issue.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 32/100