[misc] Accelerate CI for zero and dist optim (#5758)

* remove fp16 from lamb

* remove d2h copy in checking states

---------

Co-authored-by: Edenzzzz <wtan45@wisc.edu>

colossalai/testing/comparison.py

@@ -22,8 +22,6 @@ def assert_close_loose(a: Tensor, b: Tensor, rtol: float = 1e-3, atol: float = 1
         b,
         rtol=rtol,
         atol=atol,
-        msg=f"Tensor not close, shape: {a.shape} vs {b.shape}, \
-                                                dtype: {a.dtype} vs {b.dtype}",
     )