refactor: replace torch.no_grad with torch.inference_mode (where possible)

[anakin87]

Related Issues

None

Proposed Changes:

As explained in this PyTorch tweet, replacing torch.no_grad with torch.inference_mode can lead to performance improvements.

Just a first draft to run the CI and see if anything breaks...

Checklist

• I have read the contributors guidelines and the code of conduct
• I have updated the related issue with new insights and changes
• I added tests that demonstrate the correct behavior of the change
• I've used the conventional commit convention for my PR title
• I documented my code
• I ran pre-commit hooks and fixed any issue

[anakin87]

Now the CI is passing but I'm not sure that all the modified methods are covered by tests.

@julian-risch @bogdankostic Do you think that this change is safe? What is your impression?

[julian-risch]

@anakin87 The only file that needs manual testing is haystack/utils/augment_squad.py. Maybe you could run that with and without your changes to compare the results? I don't think that your changes will break anything and I am confident that it will give some speed improvements. So I am very much looking forward to seeing this PR merged. Good work! 👍
A bonus would be to have some speed measurements.

[anakin87]

Hey @julian-risch, thanks for your feedback!

Test augment_squad

In order to properly test haystack/utils/augment_squad.py with inference mode, I ran this Colab notebook.
The result is fine.

Speed measurements

I also thought to use the same notebook to record some time measurements and compare them with the original no_grad (this other notebook).

In particular, I executed the following cell:

%%timeit -n 1 -r 5
!python augment_squad.py \
    --squad_path data/distil_a_reader/squad/squad_small.json \
    --glove_path data/distil_a_reader/glove/glove.6B.300d.txt \
    --output_path augmented_dataset.json \
    --multiplication_factor 2

I obtained:

• no_grad : 2min 27s ± 4.67 s per loop (mean ± std. dev. of 5 runs, 1 loop each)
• inference_mode: 2min 12s ± 2.91 s per loop (mean ± std. dev. of 5 runs, 1 loop each)

The results would sound exciting but I'd take them with a grain of salt. In fact, I've found very few resources that measure the impact of this change and some (like this, on a different domain) report small performance improvements.

In any case, PyTorch docs say this:

If it works out of the box for your use case it’s a free performance win.

[julian-risch]

@anakin87 That's great news! Thank you for the extra effort of conducting speed experiments. I have one last question before approving and merging this PR. There are two calls to torch.no_grad() in haystack/nodes/reader/table.py and two in haystack/modeling/training/base.py. Is there a specific reason why not to use inference_mode there instead of no_grad too? Or can we also replace these four calls? It looks like the latter to me, no? Either way this is a great PR! 👍

[anakin87]

There are two calls to torch.no_grad() in haystack/nodes/reader/table.py and two in haystack/modeling/training/base.py. Is there a specific reason why not to use inference_mode there instead of no_grad too?

Yes. At first i tried to apply the change everywhere.
As you can see from the failed tests on the first few commits, i ran into errors (RuntimeError: Inference tensors cannot be saved for backward. To work around you can make a clone to get a normal tensor and use it in autograd.) and reverted the changes that triggered these errors.

[julian-risch]

These changes look good to me! 👍 Thanks again for testing the augment_squad.py script and comparing the speed before and after the changes.

[sjrl]

There are two calls to torch.no_grad() in haystack/nodes/reader/table.py and two in haystack/modeling/training/base.py. Is there a specific reason why not to use inference_mode there instead of no_grad too?

Yes. At first i tried to apply the change everywhere. As you can see from the failed tests on the first few commits, i ran into errors (RuntimeError: Inference tensors cannot be saved for backward. To work around you can make a clone to get a normal tensor and use it in autograd.) and reverted the changes that triggered these errors.

@julian-risch I've been working on the TableQA code for a client project. I also saw this error while working there and I believe I can fix it with minimal changes. Would you like me to open a new PR for that?

Update: Ahh never mind. I remember now this would require an update in the transformers library since the error results from there.

[julian-risch]

@sjrl Thanks to your PR Haystack's transformers dependency has been updated. Do you know whether that allows us to use inference_mode now also where we previously got the error mentioned above?

[sjrl]

Yes it does! It also requires a small code change in our code base which I've opened in PR #3731