Feature/double precision

[ethanwharris]

What does this PR do?

Fixes #2497

Before submitting

• Was this discussed/approved via a GitHub issue? (not for typos and docs)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure your PR does only one thing, instead of bundling different changes together?
• Did you make sure to update the documentation with your changes? (if necessary)
• Did you write any new necessary tests? (not for typos and docs)
• Did you verify new and existing tests pass locally with your changes?
• Did you update the CHANGELOG? (not for typos, docs, test updates, or internal minor changes/refactorings)

PR review

Anyone in the community is free to review the PR once the tests have passed.
Before you start reviewing make sure you have read Review guidelines. In short, see the following bullet-list:

• Is this pull request ready for review? (if not, please submit in draft mode)
• Check that all items from Before submitting are resolved
• Make sure the title is self-explanatory and the description concisely explains the PR
• Add labels and milestones (and optionally projects) to the PR so it can be classified

Did you have fun?

Make sure you had fun coding 🙃

[justusschock]

Hi @ethanwharris ,

Great to see you again :)

And great PR!
Just some smaller comments :)

[ethanwharris]

Hi @justusschock,

Thanks for the great suggestions, all done :)

[williamFalcon]

omg.... i love this feature
😍😍😍😍😍

[rohitgr7]

can we add support with deepspeed too?
https://github.com/PyTorchLightning/pytorch-lightning/blob/3a56a6024e0d8b239801cce558381807c24ba3d0/pytorch_lightning/plugins/training_type/deepspeed.py#L46-L48

not sure why it's handled this way in here, not an expert with deepspeed.

[SeanNaren]

Thanks for the awesome contribution @ethanwharris!

Out of curiosity, if the user doesn't use the forward function of the LightningModule and instead calls independent modules, does it still work? i.e if I define my training_step like:

 def training_step(self, batch, batch_idx):
        output = self.layer((batch, torch.ones_like(batch).long())) # use layer instead of self(...)

[tchaton]

Nice !

[ethanwharris]

@SeanNaren That's an excellent point, totally forgot about that haha.

@SeanNaren @carmocca @justusschock I've made some changes to address your comments:

• Moved patch logic to seperate class that now handles teardown (much cleaner than before)
• Added patch to training_step etc. methods so that it all still works if the model is called directly without forward (and added a test of this behaviour)

@rohitgr7 Not sure about deepspeed, don't feel comfortable trying to add it as I don't currently have a way of testing it out. It seems that at least DeepSpeed ZeRO requires precision=16:
https://github.com/PyTorchLightning/pytorch-lightning/blob/853523ee643fe0f0cc30d40d9e85a8869e7edfd8/tests/plugins/test_deepspeed_plugin.py#L193

Should all be ready for review now 😃

[carmocca]

LGTM. Can you have a look at the failing tests?

[justusschock]

2 more minor comments, which allow subclassing (and don't hardcode the class name, else LGTM

[ethanwharris]

@carmocca I have finally fixed the tests 😩. trainer.predict calls torch.set_grad_enabled(False) which was then killing gradients for subsequent tests. I've added a call to torch.set_grad_enable(True) which fixes this for now. I'll open an issue to change the predict behaviour to re-enable grads, maybe just switch to with torch.no_grad.