[PyTorch] Debug amax reductions in eval mode and async amax reductions

[timmoon10]

This PR fixes two bugs related to amax reductions:

1. In the forward pass, we only launch an amax reduction if a module is in training mode. However, we always launch amax reductions in the backward pass. This causes runtime errors in the backward pass of modules in evaluation mode, e.g. LoRA frozen layers. Confusing no_grad and evaluation mode seems to be a common mistake in PyTorch. This PR fixes this by checking if a module is in training mode in its backward pass, similar to how we do it in the forward pass.
2. When asynchronous amax reductions are enabled, we currently sync the reduction after the first module's amax and scale update. I would appreciate sanity checking, since I would expect this to have caused non-deterministic numerical errors in the scale updates. This PR avoids this by making sure the async reduction is finished before any amax and scale updates. See Async amax reduction #118 (comment).

I've attempted to keep this PR small since #575 touches a lot of the amax reduction logic. In the future, I think it would be worthwhile reworking the async amax reductions since it currently doesn't do much overlapping (it is launched when entering fp8_autocast and synchronized before the first TE module's forward).

[timmoon10]

/te-ci pytorch

[ksivaman]

Note: #575 already revokes the async amax reduction and addresses these fixes (including overhauling the current system for amax reduction/update), and given these would land in the same release, shall we close this? @timmoon10

[timmoon10]

/te-ci pytorch