lower precision setting

Author: inkcherry

tests/unit/model_parallelism/test_autotp_training.py

@@ -216,7 +216,7 @@ def testRowParallel(self, tp_size: int, overlap_comm: bool):
 
         torch_grad = torch.chunk(torch_linear.weight.grad, tp_size, dim=1)[groups.get_tensor_model_parallel_rank()]
         assert torch.allclose(linear.weight.grad, torch_grad.to(get_accelerator().current_device()), atol=1e-3)
-        assert torch.allclose(out, torch_out.to(get_accelerator().current_device()), atol=1e-3)
+        assert torch.allclose(out, torch_out.to(get_accelerator().current_device()), atol=1e-2)
 
     def testColumnParallel(self, tp_size: int, overlap_comm: bool):
         skip_on_device()
@@ -269,7 +269,7 @@ def testColumnParallel(self, tp_size: int, overlap_comm: bool):
         assert torch.allclose(linear.weight.grad, torch_grad.to(get_accelerator().current_device()), atol=1e-3)
         assert torch.allclose(cur_device_out.to(get_accelerator().current_device()).contiguous(),
                               out.contiguous(),
-                              atol=1e-3)
+                              atol=1e-2)
 
 
 @pytest.mark.sequential