Slow featurization when running inference for same input with multiple seeds

[y1zhou]

In run_alphafold.py, the call stack is main -> process_fold_input -> predict_structure -> featurisation.featurise_input, which calls data_pipeline.process_item repeatedly on the same inputs with different seeds. Within WholePdbPipeline.process_structure (the workhorse of process_item), the random seed is only used once towards the end of the method in features.RefStructure.compute_features. I'm wondering if the rest of the process_structure method could be taken outside of the for-loop, as it seems to be generating the same features anyways? Happy to work on a draft PR if this is the case. Thanks!

[Augustin-Zidek]

Hi, thanks for reporting this. I agree, this would be a nice feature. Could this be achieved with a well-placed @functools.lru_cache decorator? A draft PR tested with your setup would help.

I am also wondering what the performance bottleneck is in featurisation, it should typically be very cheap. Is your input sequence very long, or has many ligands, or with very deep MSA?

[y1zhou]

Sounds good I'll take a shot at this and run some local benchmarks to see if things improve on my end.

Regarding the performance bottleneck, we have a setup where the script runs in pods with fractions of CPUs and a full GPU assigned, which I'm guessing is why the featurization is slow here. We split the MSA search data pipeline and inference pipelines, so it was a bit of a surprise when we saw the GPUs sitting idle for a long time before the GPU usage comes in. On a test case with 872 tokens we observe ~10s per seed for featurization, and it's all protein chains with regular-depth MSA from the AF3 data pipeline.

[Augustin-Zidek]

Thanks for the extra details! It would be interesting to profile the run to see where the bottleneck is on the featurisation side if this is something you can do in your setup.

[y1zhou]

PR submitted! I'll try to run some py-spy profiling over the weekend to see if the featurization bottleneck can be spotted, but the cached version should already save a significant amount of time when there's many seeds.

[y1zhou]

Here's a quick profile for running the test case with four seeds: benchmark-af3-nocache.svg vs. benchmark-af3-lru-cache.svg.

1. As expected, the featurization cost is much cheaper than predict_structure. However with the cache we still reduce the number of samples in featurization from 2138/23460=9.1% to 1238/22505=5.5%. With a lot more seeds the % saved would shrink though, since GPU inference time would dominate.
2. Within process_structure, in nocache.svg (current main branch) the most expensive steps are MSA featurization, batch_ref_structure construction, and deterministic_ref_structure construction. With the LRU cache the first two steps are only computed once, so if MSAs are very deep I would expect the later seeds to be much faster.
3. In lru-cache.svg deterministic_ref_structure construction (pipeline.py:519, features.py:1720) has ~3.9x profiled samples vs. batch_ref_structure construction (same line 1720 under _process_structure_seed_independent), confirming that we are hitting the cache in seeds 2-4.