[RFC] Require bitwise equivalence for SimpleFSDP numerics (#1743)

This PR appears to pass. I am not sure if it's wise though.

Signed-off-by: Edward Z. Yang <[email protected]>

Author: ezyang

torchtitan/experiments/simple_fsdp/tests/test_numerics.py

@@ -93,7 +93,7 @@ def test_replicate_convergence(self):
         for fsdp2_loss, simple_fsdp_replicate_loss in zip(
             fsdp2_losses, simple_fsdp_replicate_losses
         ):
-            assert torch.allclose(fsdp2_loss, simple_fsdp_replicate_loss)
+            assert torch.equal(fsdp2_loss, simple_fsdp_replicate_loss)
 
     def test_fullyshard_convergence(self):
         # unit test for fully_shard mode
@@ -109,7 +109,7 @@ def test_fullyshard_convergence(self):
         for fsdp2_loss, simple_fsdp_fullyshard_loss in zip(
             fsdp2_losses, simple_fsdp_fullyshard_losses
         ):
-            assert torch.allclose(fsdp2_loss, simple_fsdp_fullyshard_loss)
+            assert torch.equal(fsdp2_loss, simple_fsdp_fullyshard_loss)
 
     def test_hybridshard_convergence(self):
         # unit test for hybrid_shard mode
@@ -125,4 +125,4 @@ def test_hybridshard_convergence(self):
         for fsdp2_loss, simple_fsdp_hybridshard_loss in zip(
             fsdp2_losses, simple_fsdp_hybridshard_losses
         ):
-            assert torch.allclose(fsdp2_loss, simple_fsdp_hybridshard_loss)
+            assert torch.equal(fsdp2_loss, simple_fsdp_hybridshard_loss)