Fix an issue in sync_amax

[y-sq]

To fix this error

RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.cuda.FloatTensor []] is at version 1; expected version 0 instead.

---

Also tried

@torch.no_grad()
 def sync_float8_amax_and_scale_history(

which didn't work.

---

We can look into if there are any better ways to fix this.

---

Test Plan:
./test/test_fsdp.sh

[facebook-github-bot]

@y-sq has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[facebook-github-bot]

@y-sq merged this pull request in 31fba04.