Fix: no_grad with AMP bug

[baskrahmer]

Fixes #20644

Note however that this would affect performance for other users, so the question is whether it is worth optimizing for this edge case that is fundamentally a torch bug.

cc @Borda

---

📚 Documentation preview 📚: https://pytorch-lightning--20921.org.readthedocs.build/en/20921/

[Borda]

Suggested change

	return torch.autocast(
	self.device, dtype=(torch.bfloat16 if self.precision == "bf16-mixed" else torch.half), cache_enabled=False
	)
	dtype = torch.bfloat16 if self.precision == "bf16-mixed" else torch.half
	return torch.autocast(self.device, dtype=dtype, cache_enabled=False)

[Borda]

Note however that this would affect performance for other users, so the question is whether it is worth optimizing for this edge case that is fundamentally a torch bug.

Then se shall report it and offer a fix in torch
Then, if it is accepted and released, we shall have a version switch in our codebase, so newer Torch versions won't need this compared to the old one. Does it...

BTW, have you measured the performance drop?
cc: @lantiga

[baskrahmer]

@Borda it is a long-standing issue in torch. I can try to make a fix if I have some time, but I think it could be complex.

But I agree with you that it should be fixed in torch ideally. Just wanted to open this PR to showcase what a workaround on our end would look like. Shall I close it?

I haven't measured the performance drop since it will vary strongly across architectures and probably also hardware setups.